Adaptive optical correction for in vivo two-photon fluorescence microscopy with neural fields

Abstract

Adaptive optics (AO) restore ideal imaging performance in complex samples by measuring and correcting optical aberrations, but often require custom-built microscopes with carefully aligned wavefront sensing/shaping devices and can be susceptible to sample motion. Here we describe NeAT, a computational framework using neural fields for AO two-photon fluorescence microscopy. NeAT estimates wavefront aberration and recovers sample structure from a 3D image stack without requiring external datasets for training. Incorporating motion correction in learning and correcting conjugation errors commonly found in commercial microscopes, NeAT is designed for deployment in biological laboratories for in vivo imaging. We validate NeAT’s performance using a custom-built microscope with a wavefront sensor under varying signal-to-noise ratios, aberration, and motion conditions. With a commercial microscope, we demonstrate real-time aberration correction for in vivo morphological and functional imaging in the living mouse brain, with NeAT improving signal and accuracy of glutamate and calcium imaging of synapses and neurons.

Introduction

Fluorescence imaging of living biological organisms provides mechanistic insights into their physiology. Two-photon (2P) fluorescence microscopy is an essential tool for live imaging, probing structure and function at subcellular resolution deep within complex tissues¹. However, as 2P excitation light propagates through tissue, its wavefront accumulates optical aberrations from refractive index mismatches, leading to reduced fluorescence signal, resolution, and contrast. When these sample-induced aberrations are measured and corrected, the excitation light can form a diffraction-limited focus, increasing fluorescence signal and improving the accuracy of structural and functional characterization.

Adaptive optics (AO)^2–6 measure aberration and correct it with wavefront-shaping devices, such as deformable mirrors (DM) and liquid-crystal spatial light modulators (SLM). AO methods can be grouped into direct wavefront sensing methods, which use a wavefront sensor for aberration measurement, and indirect wavefront sensing methods, which include approaches utilizing machine learning for wavefront estimation^7–12.

Regardless of aberration measurement scheme, AO methods are generally developed for and deployed on custom-built microscopes, in which individual optical components are carefully conjugated and aligned to ensure optimal imaging and correction performance. However, microscopes in a general laboratory setting often have imperfect conjugation and misalignment of optical components, with commercial microscopes additionally suffering from limited access and adjustability of their optical paths. Furthermore, sample motion during in vivo imaging leads to artifacts that degrade aberration measurement accuracy, a problem that can be particularly severe for deep tissue imaging as well as for indirect wavefront sensing methods that utilize serial measurement of images and signals⁴.

Here, we describe NeAT, Neural fields for Adaptive optical Two-photon fluorescence microscopy. It utilizes neural fields to represent a sample’s three-dimensional (3D) structure and incorporates computational architectures to enhance AO performance for imperfect microscopes and living samples. By incorporating an image-formation model for two-photon fluorescence microscopy that accounts for both aberration and sample motion as a physics prior, NeAT accurately estimates aberration from a single fluorescence image stack without requiring external datasets for training, even in the presence of motion artifacts. NeAT also corrects for conjugation errors in the microscope system, ensuring that the corrective phase pattern displayed on a wavefront-shaping device accurately cancels out aberration after propagation through imperfectly conjugated and misaligned optics. Lastly, NeAT jointly recovers the sample’s 3D structure with aberration. In scenarios where additional imaging with aberration correction is not needed, NeAT eliminates the need for corrective devices, further reducing system cost and complexity.

The paper is structured as follows. First, we implement NeAT in a perfectly conjugated 2P fluorescence microscope equipped with a wavefront sensor for direct wavefront sensing (DWS). We then compare the performance of NeAT with ground-truth aberration measurement by DWS both in vitro and in vivo, and determine its performance limits in terms of signal-to-noise ratio (SNR), aberration severity, and sample motion. Next, we implement NeAT in a commercial microscope with imperfect conjugation and evaluate its real-time aberration correction performance for in vivo morphological and functional imaging within the living mouse brain.

Results

NeAT, a general-purpose AO framework in 2P fluorescence microscopy using neural fields

NeAT is designed to jointly estimate wavefront aberration and recover sample structure from an input 3D 2P fluorescence image stack (Fig. 1). It utilizes neural fields – implicit functions represented by a coordinate-based neural network across spatial coordinates¹⁴ – to represent sample structure (Fig. S1a). NeAT also incorporates a mathematical image-formation model for 2P fluorescence microscopy into the learning process, which involves aberration and structural estimation, as well as motion correction through learnable image transformations. During the learning process, NeAT aims to reproduce an image stack closely resembling the input by iteratively adjusting its parameters. This process requires no external supervision.

Figure 1. NeAT estimates aberration and recovers structure from a 3D input image stack.

(a) Schematic of NeAT’s function and integration into an adaptive optics (AO) imaging pipeline. Lateral (xy) and axial (xz) maximum intensity projections (MIPs) of a 3D two-photon (2P) fluorescence image stack (g) as the input to NeAT. The learning process of NeAT corrects for sample motion, and outputs an estimated aberration (Zernike coefficients ${\alpha }_n^m$) and sample structure (s). Aberration output is then used to calculate a corrective wavefront to be applied to a deformable mirror (DM) or liquid-crystal spatial light modulator (SLM) for aberration correction. If present, microscope conjugation errors are measured by NeAT and compensated for before applying corrective wavefront. Scale bar: 5 μm. (b) If present, sample motion artifacts in g are corrected by applying learnable transformations A. (c) NeAT optimizes network weights ($\theta$), ${\alpha }_n^m$, and A to minimize loss function $ℒ\left(Ag,\widehat{g}\right)+ℛ\left(s\right)$ that compares the input image stack $Ag$ with an image stack $\widehat{g}$. $\widehat{g}$ is computed from a structure represented by a neural field ($s=f_{\theta }\left(\mathit{r}\right)$) and a 3D point spread function calculated from ${\alpha }_n^m$ (PSF h; axial MIP shown). (d) If present, conjugation errors are estimated as $\widehat{H}$ and compensated for by applying ${\widehat{H}}^{-1}$ to the corrective phase pattern ($\widehat{\phi }$). Scale bars: 5 μm.

The input for NeAT is an image stack (g) acquired through scanning (Fig. 1a). Artifacts caused by sample motion (e.g., body movement, breathing, and heartbeat) in the z stack, if present, are corrected by a set of affine transformations (A) whose parameters are optimized during the learning process (Fig. 1b, Fig. S1b). In the absence of motion artifacts, A is set as an identity operator and excluded from the learnable parameters.

The image-formation model consists of three components: point spread function (PSF, h), structure (s), and baseline (b) (Fig. 1c). The PSF h is computed as¹⁵:

$$h\left(\mathit{r};\alpha \right)={\left|ℱ\left\{P\left(u,v\right)e^{i\phi \left(u,v;\alpha \right)}{e}^{-2\pi iz\sqrt{k_0^2-u^2-v^2}}\right\}\right|}^4.$$ (1)

Here, r represents the spatial coordinates $\left(x,y,z\right)$ near the focal plane. $P\left(u,v\right)$ and $\phi \left(u,v;\alpha \right)$ stand for the amplitude and phase maps in the coordinates $\left(u,v\right)$ within the circular pupil of the objective lens, respectively. $\phi \left(u,v;\alpha \right)$ is a linear combination of Zernike modes with their associated coefficients $\alpha$, with $\phi \left(u,v;\alpha \right)=\sum _{\left|m\right|\le 3}\sum _{n=2}^4{\alpha }_n^m{Z}_n^m\left(u,v\right)$. Here m and n stand for the angular meridional frequency and radial order, respectively, following the American National Standards Institute (ANSI) standard convention for Zernike modes. $\alpha$, a 1D tensor, is a set of learned Zernike coefficients (Fig. S1c). We constrain our aberration estimation to up to fourth-order Zernike modes, excluding tip, tilt, defocus, and quadrafoil (i.e., $Z_n^m$ with $2\le n\le 4$ and $\left|m\right|\le 3$), based on prior studies^16–20. Tip, tilt, and defocus do not affect 2PFM image quality. Quadrafoil is excluded as its inclusion often yields inaccurate estimations under low-signal in vivo imaging conditions.

3D structure s is rendered by a neural field (Fig. 1c, Fig. S1a). It takes the spatial coordinates r as input and involves both Fourier-domain spatial encoding^21,22 and a multi-layer perceptron^13,23. This formulation follows the original neural field framework¹³, where spatial coordinates are mapped to radial Fourier features²¹ (i.e., $\gamma \left(\mathit{r}\right)={\left[\sin \left(2^l{\mathit{R}}_{\theta }\mathit{r}\right),\cos \left(2^l{\mathit{R}}_{\theta }\mathit{r}\right)\right]}_{l=0,\cdots ,L-1}$, where ${\mathit{R}}_{\theta }={\left\{\left[\begin{array}{cc}\cos \left({\theta }_k\right)& -\sin \left({\theta }_k\right)\\ \sin \left({\theta }_k\right)& \cos \left({\theta }_k\right)\end{array}\right]\right\}}_{k=1}^K$), $\mathit{r}=\left(x,y\right)$ denotes the spatial coordinates, L controls the maximum radial frequency depth, and K determines the number of angular samples over $2\pi$. The encoded features are then passed through a multilayer perceptron $f_{\theta }$, which represents the underlying signal as a continuous function (Fig. S1a). s is parametrized by the network weights $\theta$ and expressed as $f_{\theta }\left(\gamma \left(\mathit{r}\right)\right)$.

The baseline term $b\left(\mathit{r}\right)$ is modeled as the multiplication of three 2D tensors that represent baseline elements along each of the x, y, and z axes (Fig. 1c, Fig. S1d). This term accounts for both the offset due to background fluorescence and noise and, if present, signal decrease along the z axis due to scattering and absorption by tissue.

The image-formation model computes an image stack $\widehat{g}$ from PSF $h\left(\mathit{r};\alpha \right)$, structure $s=f_{\theta }\left(\gamma \left(\mathit{r}\right)\right)$, and baseline $b\left(\mathit{r}\right)$ by convolving PSF with structure before summation with baseline:

$$\widehat{g}\left(\mathit{r}\right)=f_{\theta }\left(\gamma \left(\mathit{r}\right)\right)⊛h\left(\mathit{r};\alpha \right)+b\left(\mathit{r}\right).$$ (2)

NeAT then compares the input stack g (or more generally with motion correction, $Ag$) and the computed stack $\widehat{g}$. It runs an optimization process to update the learnable parameters over iterations to minimize the loss function (Fig. 1c, Fig. S1e):

$$\underset{\theta ,\alpha ,A,b}{\text{min}}\left(ℒ\left(Ag,\widehat{g}\right)+ℛ\left(s\right)\right).$$ (3)

The fidelity term $ℒ\left(Ag,\widehat{g}\right)$ is a weighted sum of SSIM (Structural Similarity Index Metric)²⁴ and rMSE (relative Mean-Squared Error)^25–27 between the two stacks. SSIM evaluates similarity between $Ag$ and $\widehat{g}$ and has been widely used as both an image quality metric^24,28–33 and a loss function in computational imaging^34–36. The rMSE term computes a weighted L2 loss that reduces the influence of bright pixels and places greater emphasis on minimizing errors in dark regions. The regularization term $ℛ\left(s\right)$ incorporates a generic prior on the spatial piecewise smoothness of the structure and is the summation of three regularizations based on second-order total variation^37,38, L1, and nonlinear diffusion³⁹. Second-order total variation and L1 regularizations are chosen for rendering spatially sparse structural features (e.g., sparsely labeled neurons). Nonlinear diffusion regularization is employed to avoid both low-frequency and high-frequency artifacts in the structure recovered by NeAT. More detailed information about the image-formation model, loss function, regularization, and two-step learning is available in Methods.

Performance validation with DWS-AO

To evaluate the accuracy of NeAT’s aberration estimation, we compared the aberration output by NeAT with the ground-truth aberration from DWS with a Shack-Hartmann wavefront sensor of fluorescence from 2P-excited guide stars^40,41, using a custom-built 2P microscope with perfect conjugation between optics, including between the X and Y galvos (Fig. S2a). Microscope system aberration was measured with DWS and corrected by a DM prior to all experiments.

We first validated NeAT’s performance using 2P imaging of fixed Thy1-GFP line M mouse brain slices. A #1.5 coverslip was placed above a brain slice at a 3° tilt, which introduced aberrations similar to those typically induced by a cranial window during in vivo mouse brain imaging⁴². We set the correction collar of the objective lens to 0.17, the nominal thickness of the coverslip. From an input image stack (Fig. 2a), NeAT output 3D neuronal structures whose lateral (xy) and axial (xz) maximal intensity projections (MIPs) showed neuronal processes as well as synaptic structures such as boutons and dendritic spines (Fig. 2b). The estimated aberration had a similar phase map to the DWS measurement with a root mean square (RMS) difference of 0.09 wave (Fig. 2c) and comparable coefficients in the dominant aberration modes, e.g., primary coma $Z_3^{\pm 1}$ and spherical $Z_4^0$ (Fig. 2d). Additional performance validation with DWS shows that NeAT produces aberration estimation comparable to DWS measurement, with RMS differences of less than ~0.1 waves for both beads and brain slices (Figs. S3, S4).

Figure 2. Performance characterization with direct wavefront sensing (DWS) AO in a custom-built 2P microscope.

(a) Lateral (xy) and axial (xz) maximal intensity projections (MIPs) of an input image stack to NeAT from a fixed Thy1-GFP line M mouse brain slice. (b) Lateral and axial MIPs of the 3D neuronal structure recovered and (c) aberration estimated by NeAT, as well as aberration measured by DWS. (d) Zernike coefficients of aberrations in c. (e) Lateral and axial MIPs of an input image stack acquired in vivo from a Thy1-GFP line M mouse brain, with motion artifacts visible in xz. (f,g) Lateral axial MIPs of the structures recovered by NeAT without (f) and with (g) motion correction. (h) Aberrations estimated by NeAT without and with motion correction, respectively, and measured by DWS. (i) Zernike coefficients for aberrations in h. (j,k) Performance versus SNR using 1-μm-diameter beads under primary astigmatism (j) and primary coma (k). PCC: Pearson correlation coefficient between recovered structures; WFE: wavefront error. (l) Performance versus aberration severity evaluated using 1-μm-diameter beads. (m,n) Performance versus SNR using a fixed mouse brain slice under primary coma (m) and secondary astigmatism (n). (o) Performance versus aberration severity evaluated using brain slice. Red dashed lines: Cutoff SNRs (j,k,m,n) and cutoff aberration RMS (l,o) from piecewise linear fits (green and blue lines). (l,o) Data are presented as mean values +/− s.e.m. (N = 3). Scale bars: 5 μm.

Next, we applied NeAT to in vivo 2P imaging of the mouse cortex. In one mouse, breathing caused lateral shifts between images at different z (Fig. 2e). Without correcting for sample motion during the acquisition of the input image stack, the algorithm misinterpreted the laterally displaced images of the same structure at different z as separate structures, leading to striated appearance in the axial MIP of its structural output (Fig. 2f). NeAT addressed this by using affine transformations A to register the image stack, with the transformation matrices jointly learned alongside other parameters (Eq. 3). With sample motion corrected, the structural output was free of striation artifacts (Fig. 2g), and the aberration output much more closely resembled the ground truth (an RMS error of 0.07 wave) than the output without motion correction (an RMS error of 0.16 wave) (Fig. 2h,i).

The effectiveness of sample motion correction depends on the SNR of fluorescence images (Methods) and the magnitude of sample motion (Fig. S5). For high SNR images (e.g., SNR of 12), NeAT could handle sample motions of ±1 μm of maximum displacement. For noisier images (e.g., SNR of 3), its accuracy decreased and could only handle sample motions with ±0.25 μm displacement. This finding offers practical guidance for optimizing surgical preparation or controlling anesthesia level to minimize sample motion during image acquisition for AO, particularly during deep tissue imaging when SNR is low.

Performance limit characterizations

After validating NeAT’s performance both in vitro and in vivo, we evaluated how robustly it performed at varying SNR levels. We varied post-objective power and acquired image stacks of 1-μm-diameter fluorescence beads at different SNRs with primary astigmatism ($Z_2^{-2}$, Fig. S6a) or primary coma ($Z_3^{-1}$, Fig. S6b) introduced to the DM. At low SNR levels (e.g., SNR < 1.5), fluorescent beads were poorly visualized and the structures output by NeAT appeared fragmented as they were fitted to noise. Only at sufficiently high SNRs did the structure resemble beads (Fig. S6a, b). We quantitatively evaluated NeAT’s performance to identify the cutoff SNR below which NeAT’s performance deteriorated abruptly³⁵. We computed the Pearson correlation coefficient (PCC) between the recovered structures at different SNRs and that from an image stack acquired with no aberration and high SNR (SNR > 7, Fig. S6a, b). By fitting the PCC values to a piecewise linear curve with two distinct slopes, we identified the cutoff SNR as 1.51 for astigmatism (Fig. 2j) and 1.60 for coma (Fig. 2k). Below the cutoff SNRs, the accuracy of structural recovery decreased, as indicated by an abrupt drop of PCC values (blue curve, Fig. 2j,k); The accuracy of aberration estimation also decreased, as indicated by an increase in wavefront error (quantified by the RMS error between NeAT’s estimate and ground-truth aberrations; green curve, Fig. 2j,k).

We repeated the experiment on a fixed Thy1-GFP line M mouse brain slice to determine whether similar limits applied to spatially extended biological structures. In this case, we applied primary coma ($Z_3^{-1}$, Fig. S6c) and secondary astigmatism ($Z_4^{-2}$, Fig. S6d) to the DM separately. Similarly to beads, low-SNR images were associated with structures dominated by artifacts (Fig. S6c, d). As before, we calculated the PCC between the recovered structures at different SNRs and the ground truth from an image stack acquired with no aberration and high SNR (SNR > 5, Fig. S6c, d). We found that the cutoff SNR was 1.92 for coma (Fig. 2m) and 1.52 for astigmatism (Fig. 2n), similar to the cutoff SNRs from the bead data. This suggests that at sufficiently high SNRs (SNR ≳ 3 for aberrations tested here), NeAT achieves accurate structural recovery, independent of feature characteristics.

Moreover, we characterized NeAT’s performance limit in terms of aberration severity. We randomly generated Zernike coefficients to obtain mixed-mode aberrations with RMS values ranging from 0.05 to 0.65 waves. We then applied each aberration to the DM and acquired images of beads and brain slices at SNR > 8. With the increase in aberration, fluorescence images became more degraded in resolution and contrast (Fig. S7). At the largest aberrations tested (e.g., 0.65 waves for beads and 0.43 waves for brain slices), the recovered structures no longer accurately represented the features of the beads or neurons. We computed the PCC between the structures retrieved by NeAT from images with varying levels of external aberration and the structure from an image stack without aberration. Similar to above, we defined the cutoff RMS as the value above which the PCC exhibited a sudden drop, as identified by fitting the PCC values to a piecewise linear curve with two distinct slopes. We found a cutoff RMS of 0.47 wave for 1-μm beads (Fig. 2n) and 0.30 wave for the brain slice (Fig. 2o), respectively. This difference in cutoff RMS values is expected as 3D extended structures generally pose greater challenges than beads.

Lastly, we characterized NeAT’s performance limit in terms of sampling rate by varying the pixel sizes of input image stacks. We downsampled both in vitro and in vivo image stacks of neurons by different factors to vary the input pixel size along the lateral (dx, dy) and axial (dz) axes, and compared NeAT’s performance in structural recovery and aberration estimation (Fig. S8, 9). When pixel size exceeded the Nyquist sampling criterion, the structure outputs from NeAT became inaccurate. The aberration estimation also deviated from the ground truth measured by DWS, with the estimated aberration matching the DWS measurement until lateral pixel size exceeded 0.20 μm and axial pixel size exceeded 0.67 μm, values dictated by the Nyquist condition, for both in vitro (Fig. S8a, Fig. S9a) and in vivo (Fig. S8b, Fig. S9a) cases.

NeAT corrects for conjugation errors in a commercial microscope

Having demonstrated the successful application of NeAT in a custom-built 2P fluorescence microscope and acquired a thorough understanding of its performance in relation to SNR, motion, aberration severity, and input pixel size, we next tested whether NeAT worked on a commercial 2P microscope. This was motivated by the desire to expand the application of AO beyond optical specialists to a general laboratory setting with microscopes having imperfect conjugation and misalignment of optical components, as well as limited access and adjustability of their optical paths.

We added a liquid-crystal SLM to the beam path between an excitation laser and a commercial 2P fluorescence microscope (Bergamo II, Thorlabs) (Fig. S2b). This system differs from our custom-built microscope in several ways. First, the DM, x galvo, and y galvo of the custom-built system were conjugated with pairs of lenses (Fig. S2a) to ensure that the corrective phase pattern displayed on the DM was accurately relayed to the back focal plane (BFP) of the objective lens and stayed stationary during beam scanning. But the commercial microscope, as typical for microscopes in biological laboratories, did not conjugate the two galvos but placed them close to each other. Second, while the optics of the custom-built system were carefully arranged and aligned to ensure the registration between the x and y axes of the SLM surface and the fluorescence images, the commercial microscope had multiple mirrors in an enclosed optical path whose placement and alignment were preset and not adjustable. Finally, the commercial system was designed to have the whole microscope body move in 3D to accommodate large samples, which causes alignment errors between the SLM on the optical table and the objective lens in the microscope that for heavily shared microscopes can vary daily. As a result, a wavefront applied to SLM is translated, rotated, scaled, and/or sheared at the BFP of the objective lens of the commercial microscope, which in turn degrades the performance of aberration correction.

To address this problem, we incorporated into NeAT a process to estimate and correct conjugation errors of commercial microscopes (Fig. 1a). Corrective wavefront displayed on the SLM, ${\phi }_{\mathrm{C}\mathrm{o}\mathrm{r}\mathrm{r}}$, becomes ${\phi }_{\mathrm{B}\mathrm{F}\mathrm{P}}$ at the BFP of the objective lens, with

$${\phi }_{\mathrm{B}\mathrm{F}\mathrm{P}}=H{\phi }_{\mathrm{C}\mathrm{o}\mathrm{r}\mathrm{r}}+{\Phi }_{\mathrm{S}\mathrm{y}\mathrm{s}}$$ (4)

Here ${\Phi }_{\text{Sys}}$ represents the system aberration and H is a linear geometric transformation describing the effects of conjugation errors on ${\phi }_{\mathrm{C}\mathrm{o}\mathrm{r}\mathrm{r}}$ (Fig. 3a). We model H as an affine transformation with parameters for translational, rotational, scaling, and shear transformation (Fig. 3b). For microscopes with perfect conjugation, H = I, the identity operator (i.e. translations are 0 pixels in x and y, rotation is 0 deg, scaling is 1, and shear is 0). For microscopes with conjugation errors, the process of accounting for them requires finding the transformation H and system aberration ${\Phi }_{\text{Sys}}$.

Figure 3. NeAT corrects for conjugation errors in a commercial microscope.

(a) Conjugation errors transform corrective pattern ${\phi }_{\text{Corr}}$ on SLM to ${\phi }_{\text{BFP}}=H{\phi }_{\text{Corr}}+{\Phi }_{\text{Sys}}$ at objective lens back focal plane. ${\Phi }_{\text{Sys}}$: system aberration. (b) H (with example affine parameters) translates, rotates, scales, and shears a unit square (black dashed square) to a parallelogram (red). px: pixel. (c) H is determined from image stacks of 200-nm-diameter beads acquired with calibration aberration ${\Phi }_n$ ($n=1,2,\cdots ,5$) applied to SLM. Lateral (xy) and axial (xz) MIPs of the calibration image stacks are shown. (d) Lateral and axial MIPs of image stacks of 200-nm-diameter beads acquired without system aberration correction (‘No AO’), after one iteration of AO without (‘AO1 w/o H’) or with (‘AO1 w/H’) conjugation correction, and after two iterations of AO with conjugation correction (‘AO2 w/H’). (e) Axial signal profiles of the bead marked by yellow arrowhead in d. (f) Residual aberrations estimated by NeAT from image stacks in d. (g) Left to right: aberrations (with 0.3 wave RMS) applied to SLM, estimated aberration by NeAT without conjugation correction, and estimated aberration by NeAT with conjugation correction from bead image stacks acquired with the applied aberration. Numbers to the bottom right of estimated aberrations: difference (in wave RMS) between estimated aberration and applied aberration. Scale bars: 5 μm.

We determine system aberration ${\Phi }_{\text{Sys}}$ by inputting into NeAT an image stack of 200-nm-diameter fluorescence beads acquired with a flat phase pattern applied to the SLM. The estimated system aberration from NeAT is ${\widehat{\varphi }}_0$, with

$${\Phi }_{\text{Sys}}={\widehat{\varphi }}_0+{\epsilon }_0$$ (5)

Here ${\epsilon }_0$ represents estimation error by NeAT, which should be much smaller in RMS magnitude than ${\Phi }_{\mathrm{S}\mathrm{y}\mathrm{s}}$.

To determine H, we apply 5 calibration aberrations ${\Phi }_n$ (n = 1 to 5) including primary astigmatism ($Z_2^{-2}$ and $Z_2^2$), coma ($Z_3^{-1}$ and $Z_3^1$), and spherical aberration ($Z_4^0$), to the SLM. These calibration aberrations allow us to detect translation, scaling, rotation, and shear errors in conjugation. At the objective lens BFP, these aberrations became $H{\Phi }_n+{\Phi }_{\mathrm{S}\mathrm{y}\mathrm{s}}$. With image stacks of 200-nm fluorescence beads acquired under these external aberrations as inputs (Fig. 3c), NeAT returns ${\widehat{\varphi }}_n$ (n = 1 to 5), with

$${H{\Phi }_n+{\Phi }_{\mathrm{S}\mathrm{y}\mathrm{s}}=\widehat{\varphi }}_n+{\epsilon }_n$$ (6)

Here ${\epsilon }_n$ represents estimation error by NeAT. Subtracting (5) from (6) and assuming ${\epsilon }_n-{\epsilon }_0\approx 0$, we have

$$H{\Phi }_n\cong {\widehat{\varphi }}_n-{\widehat{\varphi }}_0,n=1,2,\dots ,5$$ (7)

Now with ${\Phi }_n$ (n = 1 to 5) known, and ${\widehat{\varphi }}_n$ and ${\widehat{\varphi }}_0$ from NeAT, we determine the parameters of H by minimizing the loss function:

$$\widehat{H}=\underset{H}{\text{argmin}}\left(\sum _{n=1}^5\left|H{\Phi }_n-\left({\widehat{\varphi }}_n-{\widehat{\varphi }}_0\right)\right|\right)$$ (8)

$\widehat{H}$, the estimate for H, describes how the conjugation errors in the system translate, rotate, scale, and shear the wavefront pattern applied to the SLM on its way to the BFP of the objective lens. To correct these errors, we then apply the inverse of $\widehat{H}$, ${\widehat{H}}^{-1}$, to the aberration estimation $\widehat{\varphi }$ from NeAT and use ${\widehat{H}}^{-1}\widehat{\varphi }$ as the corrective pattern on the SLM (Fig. 1d).

For example, to correct for system aberration of the commercial microscope, we used an image stack of 200-nm fluorescence beads as input to NeAT, which returned ${\widehat{\varphi }}_0$ as the aberration estimation. Directly applying ${\widehat{\varphi }}_0$ to the SLM increased the signal of a fluorescent bead by 1.7-fold (“AO1, w/o H”, Fig. 3d,e). By also correcting for conjugation errors, ${\widehat{H}}^{-1}{\widehat{\varphi }}_0$ increased the signal by 2.2-fold (“AO1, w/ H”, Fig. 3d,e). Using the image stack acquired with ${\widehat{H}}^{-1}{\widehat{\varphi }}_0$ as input into NeAT, we obtained the residual aberration ${{\widehat{\varphi }}_0}^{\prime }$ and applied ${\widehat{H}}^{-1}\left({\widehat{\varphi }}_0+{{\widehat{\varphi }}_0}^{\prime }\right)$ to the SLM, leading to a 3.0-fold signal gain over no aberration correction (“AO2, w/ H”, Fig. 3d,e). From the image stacks acquired with these corrective patterns, NeAT estimated the residual aberrations (Fig. 3f). Consistent with the fluorescent signal measurements, conjugation error correction substantially reduced residual aberration, with 0.14 and 0.12 wave RMS after the first and second iterations of AO correction, while the residual aberration without conjugation correction had a 0.22 wave RMS.

We further tested our approach on correcting known astigmatism, coma, and spherical aberrations introduced to the SLM. From bead image stacks acquired with these aberrations applied, NeAT returned estimated aberrations (“Estimated w/o H”, Fig. 3g), which represented the wavefront distortion at the objective BPF and substantially differed from the applied aberrations (“Applied aberration”, Fig. 3g) due to conjugation errors. Transforming the estimated aberration with ${\widehat{H}}^{-1}$, we obtained aberrations with phase maps (“Estimated w/ H”, Fig. 3g) that closely matched the given aberrations in all three cases, leading to much smaller RMS errors (astigmatism: 0.087 and 0.19 wave RMS with and without H correction; coma: 0.14 and 0.19 wave RMS with and without H correction; spherical: 0.16 and 0.23 wave RMS with and without H correction). Once characterized, the same ${\widehat{H}}^{-1}$ can be applied as long as the conjugation of the microscope remains unchanged. Below, the system aberration of the commercial microscope was always corrected for “No AO” images so that improvement by AO arose from the correction of sample-induced aberrations alone.

Real-time aberration correction for in vivo structural imaging of mouse brain

We evaluated NeAT’s capacity to improve in vivo structural imaging with the commercial microscope. We acquired an image stack of a tdTomato-expressing dendrite at 350 μm depth in the primary visual cortex (V1) of a head-fixed mouse (“No AO”, Fig. 4a) and used it as input to NeAT to the SLM the corrective wavefront from NeAT with both motion and conjugation corrections, we imaged the same dendrite and observed a marked improvement in brightness (up to 1.8× for dendritic spines), resolution, and contrast (“Full correction”, Fig. 4a).

Figure 4. Real-time aberration correction by NeAT for in vivo structural imaging.

(a) Lateral (xy) and axial (xz) MIPs of image stacks of tdTomato-expressing dendrite and dendritic spines at 350 μm depth acquired with system aberration correction only (“No AO”, used as input to NeAT), with corrective wavefront estimated by NeAT with both conjugation and motion corrections (“Full correction”), motion correction only (“Without H”), or conjugation correction only (“Without motion correction”). (b) Estimated aberrations by NeAT without and with conjugation correction. (c) Lateral signal profiles along dashed lines and axial signal profiles of spine indicated by arrowhead in a. (d) Lateral and axial MIPs of Thy1-GFP line M mouse dendrites at 280 μm depth acquired with system aberration only (“No AO”, used as input to NeAT) and with correcting sample-induced aberration by NeAT. (e) Estimated aberration by NeAT. (f) Axial profiles of dendritic spines marked by arrowheads in d. (g) Lateral and axial MIPs of neuronal processes at 500 μm depth, acquired with system aberration correction only (“No AO”), aberration correction at 280 μm (“AO_{280 μm}”, used as input to NeAT; wavefront in e), and aberration correction at 500 μm (“AO_{500 μm}”) (h) Sample-induced aberration at 500 μm. (i) Lateral profiles along dashed line and axial profiles of spines indicated by arrowheads in g. Scale bars: 5 μm.

Correcting for both sample motion and conjugation error was necessary for the observed improvement. The corrective wavefront with only motion correction but not conjugation correction substantially differed from that with full correction (Fig. 4b) and led to more modest improvements in image quality (“Without H”, Fig. 4a). Only correcting for conjugation but not motion similarly underperformed (“Without motion correction”, Fig. 4a). These trends were observed quantitatively in the lateral and axial intensity profiles of three example dendritic spines (Fig. 4c).

We investigated further whether image-registration software such as the StackReg plugin in ImageJ can work similarly well to the motion correction method integrated into the learning process of NeAT. We acquired an image stack of beads with aberration, introduced simulated motion artifacts, pre-registered the stack with StackReg, and then used the resulting image stack as input to NeAT. Although structural recovery was moderately successful for beads (Fig. S10a, b), the accuracy of aberration estimation from the pre-registered image stack was inferior to inputting the un-registered stack to NeAT directly (Fig. S10c). Similarly, pre-registration with StackReg on an image stack acquired in vivo led to a corrective wavefront with smaller brightness improvement than the corrective wavefront from motion correction by NeAT (Fig. S10d,e). This can be explained by whether motion correction considers the existence of aberration. While NeAT integrates motion correction into its learning process for aberration (Eq. 3), conventional image registration is unaware of aberrations and matches features between adjacent image planes to align them, which may inadvertently reduce or exaggerate certain aberrations (e.g., the axially curved tail of comatic aberration may be straightened by StackReg).

Having established that both conjugation and motion corrections are needed for in vivo imaging using the commercial microscope, we further tested NeAT’s performance for morphological imaging deep within the brain of a Thy1-GFP line M mouse. We first used an image stack acquired at a depth of 280 μm as input to NeAT (“No AO”, Fig. 4d) to obtain the corrective wavefront (Fig. 4e, 0.36 wave RMS,), which led to resolution improvement as well as an ~2× increase in spine brightness (“AO”, Fig. 4d,f), where AO substantially enhances the resolution and contrast of fine structures such as dendritic spines (for more examples, see Fig. S11a). We then acquired an image stack at 500 μm depth while applying to the SLM the corrective wavefront at 280 μm (“AO_{280 μm}”, Fig. 4g). Using the image stack as input to NeAT, we obtained a corrective wavefront, which was then added to the corrective wavefront at 280 μm to obtain the final corrective pattern (Fig. 4h, 0.49 wave RMS). This corrective wavefront has a larger RMS magnitude than that at 280 μm, consistent with previous observation of stronger aberrations at larger imaging depths for the mouse brain⁴³. Compared to the image stacks acquired without AO (“No AO”, Fig. 4g) and with corrective wavefront at 280 μm (“AO_{280 μm}”, Fig. 4g), images after correction at 500 μm (“AO_{500 μm}”, Fig. 4g) had the best resolution and contrast, with up to a 2.4-fold increase in brightness for dendritic and synaptic structures (Fig. 4i). Here by using the corrective wavefront at a shallower depth when acquiring the input image stack at a deeper depth, we overcame the limit on aberration severity and used NeAT to correct large aberrations experienced in deep tissue imaging.

NeAT improves in vivo glutamate imaging from the mouse brain

We next used NeAT with motion and conjugation correction to improve in vivo functional imaging in head-fixed mice. We expressed the genetically encoded glutamate indicator iGluSnFR3⁴⁴ sparsely in V1 neurons (Methods). From an image stack of dendrites at 400-μm depth (Fig. 5a), NeAT returned a corrective wavefront (Fig. 5b) that substantially increased image resolution and contrast, resulting in approximately two-fold improvement in brightness as shown by axial profiles of dendritic spines (i,ii; Fig. 5c) and resolving a dendritic spine from its nearby dendrite (iii; Fig. 5c) (for more examples, Fig. S11b).

Figure 5. Real-time aberration correction by NeAT for in vivo glutamate and calcium imaging.

(a) Lateral (xy) and axial (yz) MIPs of input stacks to NeAT (“No AO”) and stacks acquired after aberration correction by NeAT (“AO”) of dendrites expressing iGluSnFR3 at 400 μm depth in mouse V1. (b) Estimated aberration by NeAT. (c) Axial profiles of spines indicated by arrowheads and lateral profiles along dashed line in a. (d) Averages of time-lapse xy images of dendrites measured without and with AO. (e) Zoomed-in views of structures in box in d. (f) Trial-averaged signal traces (F), glutamate transient traces (ΔF/F₀), orientation tuning curves and OSI values of 4 ROIs (1-4 in e). Shade and error bars: s.e.m. (g) OS spines in d color-coded by their preferred orientations measured without and with AO. (h-k) Comparisons of basal fluorescence (${\text{F}}_0$; h), glutamate transient (ΔF/F₀; i), OSIs (j,k) of 52 OS spines out of 86 total spines before and after AO correction. Two-sided paired t-test (h-j), p < 0.001. Kolmogorov-Smirnov test (k), p < 0.001. (l) Superimposed xy images of sparse tdTomato-expressing neurons (1000-nm excitation) and dense GCaMP6s-expressing neurons (920-nm excitation) at 280 μm depth. (m) Lateral (xy) and axial (xz) MIPs of dendrites (yellow box in l) measured without (“No AO”, input to NeAT) and with AO. (n) Estimated aberration by NeAT. (o) Axial profiles for spines indicated by arrowheads and lateral profiles along dashed line in m. (p) Standard deviation of time-lapse images of GCaMP6s-expressing neurons in white box in l, acquired without and with AO. (q) Trial-averaged signal traces (F), calcium transient traces (ΔF/F₀), orientation tuning curves, and OSI values of 5 ROIs (1-5 in p). Shade and error bars: s.e.m. (r-u) Comparisons of basal fluorescence (${\text{F}}_0$; h), calcium transient (ΔF/F₀; s), OSIs (t,u) of 125 OS ROIs out of 255 somatic and neuronal structures before and after AO correction. Two-sided paired t-test, p < 0.001 (r,t), p < 0.05 (s). Kolmogorov-Smirnov test (u), p < 0.001. Scale bars: (a,d,e,g) 5 μm; (l,p) 100 μm; (m) 10 μm.

Subsequently, we presented gratings drifting in eight different directions (0°, 45°,⋯, 315°;10 repetitions) to the mouse and recorded 2D time-lapse images of dendritic spines in the same FOV as in Fig. 5a at a 60 Hz frame rate, with and without the corrective wavefront applied to the SLM. With iGluSnFR3 labeling, changes in fluorescence brightness reflected glutamate release and thus synaptic input strength at these dendritic spines. Consistent with the above result, AO increased the brightness of dendrites and spines in the averaged time-lapse image (Fig. 5d; Fig. 5e, zoomed-in views of white boxes in Fig. 5d). For four representative dendritic spines (ROI 1–4, Fig. 5e), AO correction doubled the basal intensity (F₀) of their trial-averaged fluorescent traces and led to more prominent glutamate transients with larger amplitudes (ΔF/F₀) (left and middle panels, Fig. 5f). Fitting the glutamate responses to the 8 drifting grating stimuli with a bimodal Gaussian curve⁴⁵, we obtained the orientation-tuning curves for these spines (right panels, Fig. 5f). Here AO increased the response amplitudes to the preferred grating orientations and led to a higher orientation sensitivity index (OSI) for these spines. Correcting aberration also shifted the preferred orientation of some spines (e.g., ROI 3 and 4, Fig. 5f), resulting in more similar tuning preference for neighboring spines (Fig. 5g), consistent with previous findings⁴⁶. Consistently across spine populations (52 orientation-sensitive ROIs out of 86 spines, Methods), aberration correction by NeAT significantly increased basal fluorescence ${\text{F}}_{\text{0}}$ by 1.9-fold on average (two-sided paired t-test, p < 0.001, Fig. 5h). It also increased ΔF/F₀ and OSI values as indicated by pairwise comparison (two-sided paired t-test, p < 0.001, Fig. 5i and Fig. 5j, respectively) and the cumulative OSI distributions (Kolmogorov-Smirnov test, p < 0.001, Fig. 5k).

NeAT improves in vivo calcium imaging in densely labeled brains

We further demonstrated that NeAT can also be applied to densely labeled brains, a common application scenario for in vivo calcium imaging of neuronal populations. As NeAT requires an input stack of sparse structures for aberration estimation, we used viral transduction to densely express the genetically encoded calcium indicator GCaMP6s⁴⁷ and sparsely express the red fluorescent protein tdTomato in L2/3 neurons of the mouse V1 (Methods). Because aberration estimation and correction can be performed at different excitation wavelengths without compromising correction performance (Fig. S12), we used an image stack of a tdTomato-expressing neuron (inside yellow box of Fig. 5l) acquired with 1000 nm excitation light as the input to NeAT (32×32×10 μm³ stack, “No AO”, Fig. 5m) to obtain the corrective wavefront (Fig. 5n). AO visibly improved image contrast and resolution of the tdTomato-expressing neuron (“AO”, Fig. 5m), leading to a >2× increase in intensity in both axial profiles at dendritic spines and lateral profiles across dendrites (Fig. 5o).

Next, we switched the excitation wavelength to 920 nm and acquired 2D images of GCaMP6s-expressing neurons over a 484×484 μm² FOV (green channel of Fig. 5l) at 15 Hz without and with the corrective wavefront obtained by NeAT, while presenting drifting gratings to the head-fixed mouse to evoke calcium responses. The standard deviation images of the time-lapse stacks showed greater intensity differences across time frames after aberration correction (Fig. 5p, zoomed-in views on the white boxes in Fig. 5l), indicative of larger calcium transient magnitude. Indeed, for five representative ROIs (1-5, Fig. 5p), calcium transients were more apparent and had larger magnitudes in both trial-averaged fluorescence (F) and ΔF/F₀ traces with AO (left and middle panels, Fig. 5q), leading to higher orientation selectivity indices for these structures (right panels, Fig. 5q).

Over the population of 125 orientation-selective ROIs out of 255 somatic and neuronal structures within the whole FOV, we found statistically significant differences between No AO and AO conditions for both basal fluorescence $F_0$ (two-sided paired t-test, p < 0.001, Fig. 5r) and ΔF/F₀ (p < 0.05, Fig. 5s). Here the increase in basal fluorescence was less than what we observed for glutamate imaging of dendritic spines, because aberration decreases signal brightness of smaller structures such as dendritic spines more than larger structures such as somata^43,48,49. Similar to glutamate imaging, AO increased the OSIs of neuronal structures (two-sided paired t-test, p < 0.001, Fig. 5t; for cumulative distributions of OSI, Kolmogorov-Smirnov test, p < 0.001, Fig. 5u).

Discussion

In this work, we describe NeAT, a general-purpose AO framework for aberration measurement and correction for 2P microscopy using neural fields. Neural fields refer to implicit functions represented by a coordinate-based neural network across spatial coordinates¹⁴. They have been used for various computational imaging applications^{21,35,50–54}, including Neural Radiance Fields (NeRF)¹³ for 3D scene representation. NeAT has several distinct features that set it apart from NeRF and other neural field applications (detailed in Table S1), including its incorporation of a physics-based prior specific to the 2P imaging system, its estimation and correction of sample motion and microscope conjugation errors, and its joint recovery of 3D structural information alongside aberration estimation.

Using the physics-based prior of the 2P imaging system, implemented through an image-formation model that accounts for both aberrations and sample motion, NeAT accurately estimates optical aberrations from a 2P fluorescence image stack without the need for external supervision, even in the presence of motion artifacts during live animal imaging. Importantly, this functionality eliminates the need for integrating AO capabilities into microscope control software, making it applicable to existing custom-built and commercial 2P microscopy systems in general.

Furthermore, NeAT can estimate and correct conjugation errors. Such errors are common in homebuilt and commercial microscopes used in general biology laboratories and cause the distortion of the applied corrective phase pattern at the objective lens back focal plane, leading to deteriorated AO performance. NeAT measures the impact of conjugation errors on calibration aberrations and compensates for them by preemptively transforming the corrective phase pattern before it is displayed on the wavefront-shaping device. This feature would greatly facilitate the broader adoption of AO in a wide range of microscope systems.

Additionally, NeAT simultaneously recovers 3D structural information while estimating aberration during its learning process. For applications where only structural information is needed, this unique capability eliminates the requirement of wavefront-shaping devices or the need for additional imaging with AO correction, greatly lowering system complexity and cost.

We rigorously evaluated NeAT’s performance under various conditions, such as different SNR levels, aberration severity, and motion artifacts. We established the framework’s performance limits for accurate and reliable operation, providing valuable guidelines on imaging settings when applying NeAT.

Finally, we applied NeAT to in vivo structural and functional imaging of the mouse brain by a commercial microscope, demonstrating its capability in improving image quality for demanding real-life biological applications. NeAT effectively estimates and corrects aberrations deep into the mouse brain, enabling morphological imaging of synapses at 500 μm with improved resolution and contrast. It also improves the signal and accuracy of glutamate and calcium imaging of synapses and neurons in the mouse visual cortex responding to visual stimuli.

With a single z-stack as input and a computation time of a few minutes (Table S2), NeAT’s simple implementation, robust performance, and ability to correct for motion and conjugation errors in imaging systems offer great potential for broader adoption and impact in biological research than many hardware-based AO methods (Supplementary Note 1).

Methods

Custom-built two-photon microscopy with direct wavefront sensing AO

A custom-built two-photon fluorescence microscope was equipped with a wavefront sensor for DWS and described previously^55,56 (Fig. S2a). A Ti-Sapphire laser (Chameleon Ultra II, Coherent Inc.) was tuned to 920 nm output and scanned by a pair of carefully conjugated galvos (H2105, Cambridge). Pairs of achromatic doublet lenses (L3-L8) conjugated the surfaces of galvos with a DM (PTT489, Iris AO) and the BFP of an objective lens (CFI Apo LWD ×25, 1.1 NA, 2.0 mm WD, Nikon). During imaging, 2P excited fluorescence was collected by the same objective, reflected by a dichroic mirror (D2, Di02-R785-25x36, Semrock), and detected by a GaAsP photomultiplier tube (H7422-40, Hamamatsu). For wavefront sensing, the emitted 2P fluorescence was descanned by the galvo pair, reflected by a dichroic mirror (D1, Di02-R785-25x36, Semrock), and directed to a Shack-Hartmann (SH) sensor through a pair of achromatic lenses (FL = 60, 175 mm). The SH sensor consisted of a lenslet array (Advanced Microoptic Systems GmbH) conjugated to the objective BFP and a CMOS camera (Orca Flash 4.0, Hamamatsu) positioned at the focal plane of the lenslet array.

Commercial two-photon fluorescence microscopy with an AO module

The commercially available multiphoton microscope (Bergamo II, Thorlabs) used a Ti-Sapphire laser (Chameleon Ultra II, Coherent Inc.) tuned to 920 nm or 1000nm for 2P excitation. An AO module consisted of a liquid crystal SLM (1024 × 1024, HSP1K, Meadowlark Inc.) and two pairs of relay lenses (L1-L4, FL = 200, 50, 500, and 500mm) was added to the beam path on the optical table between the laser and the microscope. The laser output had its polarization rotated by an achromatic half-wave plate (AHWP05M-980, Thorlabs) to align with the SLM polarization requirement and was expanded 15 times using two beam expanders (GBE03-B, GBE05-B, Thorlabs) to fill the active area of the SLM. The two pairs of relay lenses demagnified the laser output and conjugated the SLM surface to the non-resonant galvo surface within the galvo-resonant-galvo scanning head of the microscope. A pair of scan lenses within the Bergamo II microscope (L5-L6, FL=50 and 200mm) relayed the laser to the BFP of a water-dipping objective (25×, 1.05 NA, 2mm WD, Olympus). Fluorescence emission was collected through the objective and detected by two GaAsP photomultiplier tubes (PMT 2100, Thorlabs) for two-color imaging of green (525/50 nm emission filter) and red (607/70 nm emission filter) fluorescence, respectively.

Animals and surgical procedures

All animal experiments were conducted in accordance with the National Institutes of Health guidelines for animal research. Procedures and protocols involving mice were approved by the Institutional Animal Care and Use Committee at the University of California, Berkeley. In vivo imaging experiments were performed using 2-4-month-old wild-type (C57BL/6J) or Thy1-GFP line M mouse lines.

Cranial window and virus injection surgeries were conducted under anesthesia (2% isoflurane in O₂) following established procedures^46,57. For in vivo glutamate imaging, sparse expression of iGluSnFR3 was achieved in V1 L2/3 by injecting a 1:1 mixture of diluted AAV2/1-Syn-Cre virus (original titer 1.8 × 10¹³ GC/ml, diluted 10,000-fold in phosphate-buffered saline) and AAV-hSyn-FLEX-iGluSnFR3-v857-SGZ at multiple sites 150-250 μm below pia. 25 nl of the virus mixture was injected at each site. For in vivo calcium imaging, dense expression of GCaMP6s and sparse expression of tdTomato was achieved in V1 L2/3 by co-injecting a 1:1:1 mixture of diluted AAV2/1-Syn-Cre virus (original titer 1.8 × 10¹³ GC/ml, diluted 1,000-fold in phosphate-buffered saline), AAV2/1-CAG-FLEX-tdTomato (6 × 10¹³ GC/ml), and AAV1-Syn-GCaMP6s-WPRE-SV40 (1 × 10¹³ GC/ml) at multiple sites 150-250 μm below the pia. 25 nl of the virus mixture was injected at each site. A cranial window, made of a glass coverslip (Fisher Scientific, no. 1.5), was embedded in the craniotomy and sealed in place with Vetbond tissue adhesive (3M). A metal head post was attached to the skull using cyanoacrylate glue and dental acrylic. After 3 weeks of expression and 3 days of habituation for head fixation, in vivo imaging was conducted in head-fixed mice under anesthesia (1% isoflurane in O₂) for structural imaging and in lightly anesthetized mice (0.5% isoflurane in O₂) for functional imaging.

Loss function and regularization in self-supervised learning process

The fidelity term $ℒ\left(Ag,\widehat{g}\right)$ in the loss function (Eq. 3) is represented as a weighted sum of SSIM²⁴ and rMSE^25–27 as follows:

$$ℒ\left(Ag,\widehat{g}\right)=\gamma (1-\mathrm{S}\mathrm{S}\mathrm{I}\mathrm{M}\left(Ag,\widehat{g}\right))+\left(1-\gamma \right)\mathrm{r}\mathrm{M}\mathrm{S}\mathrm{E}(Ag,\widehat{g}),$$ (9)

where rMSE is defined as

$$\mathrm{r}\mathrm{M}\mathrm{S}\mathrm{E}\left(Ag,\widehat{g}\right)={\left(\frac{Ag-\widehat{g}}{\mathrm{s}\mathrm{g}\left(\widehat{g}\right)+{\epsilon }_l}\right)}^2.$$ (10)

Here $\mathrm{s}\mathrm{g}\left(\cdot \right)$ denotes a stop-gradient operation that treats its argument as a constant, employed for numerical stability during backpropagation²⁵. This operation prevents the denominator from being differentiated with respect to the network weights, thereby avoiding division-by-zero instabilities or exploding gradients. The parameter $\gamma$ controls the weight between the two terms. It is set to 0.25 if the RMS contrast of the image stack’s background pixels, i.e., ${\epsilon }_l={\sigma }_b\left(g_{bfr}\right)$, is larger than 0.03, where ${\sigma }_b\left(\cdot \right)$ computes the standard deviation of the background pixels of the operand. If the contrast is smaller than 0.02, $\gamma$ is set to 1.0. Otherwise, $\gamma$ is linearly interpolated between 0.25 and 1.0. Here, $g_{bfr}$ represents a background-fluctuation-removed version of $g$, introduced to remove any unwanted low-frequency fluctuations in the images that could otherwise exaggerate the standard deviation.

The regularization term $ℛ\left(s\right)$ in the loss function (Eq. 3) is designed to render spatially sparse and smooth structural details, serving as a generic prior that reflects structural features of mouse brain neurons. It includes three regularization terms: second-order total variation (TV) ${ℛ}_{tv}\left(s\right)$^37,38, L1 regularization ${ℛ}_{\mathrm{L}1}\left(s\right)$, and nonlinear diffusion (NLD) ${ℛ}_{\mathrm{N}\mathrm{L}\mathrm{D}}\left(s\right)$³⁹. Since the 3D structure $s\left(\mathit{r}\right)=f_{\theta }\left(\gamma \right(\mathit{r}\left)\right)$ is represented implicitly by a neural field, the required spatial derivatives are evaluated directly through automatic differentiation using PyTorch’s autograd⁵⁸ during the learning process. The convergence behavior of the fidelity and regularization terms remains stable across epochs (Fig. S13).

First, second-order TV ${ℛ}_{tv}\left(s\right)$ aims to recover smooth profiles from noisy measurements by sparsifying the spatial gradient components. Unlike first-order TV⁵⁹, which uses first-order derivatives, second-order TV uses second-order derivatives to avoid staircase artifacts^37,38. In our implementation, we further applied a nonlinear tone mapping function²⁵, an approximated logarithmic function (Eq. 12) that strongly penalizes errors in regions with low intensity values. For simplicity, the spatial coordinates $\left(x,y,z\right)$ are expressed as $\left(x_1,x_2,x_3\right)$ below.

$$tv\left(s\right)=\sum _{1\le i\le j\le 3}{\text{TV}}_{x_i{x}_j}\left(s\right),\text{with}{\text{TV}}_{x_i{x}_j}\left(s\right)=\left|\frac{{\partial }^{2}s}{\partial x_i\partial x_j}\right|,$$ (11)

where $tv\left(s\right)$ includes all second-order TV terms across x, y, z, xy, yz, and zx dimensions.

$${ℛ}_{tv}\left(s\right)=\mathrm{l}\mathrm{o}\mathrm{g}\left(tv\left(s\right)+{\epsilon }_{tv}\right)\simeq \frac{tv\left(s\right)}{\mathrm{s}\mathrm{g}\left(tv\right)+{\epsilon }_{tv}},$$ (12)

where $\mathrm{s}\mathrm{g}\left(\cdot \right)$ indicates the same stop-gradient operation as above, and ${\epsilon }_{tv}$ is determined from the input image stack g as the smallest standard deviation of second-order difference $tv\left(g\right)$, that is,

$${\epsilon }_{tv}=\underset{1\le i\le j\le 3}{\min }\left({\sigma }_b\left({\text{TV}}_{x_i{x}_j}\left(g\right)\right)\right).$$ (13)

Second, L1 regularization ${ℛ}_{\mathrm{L}1}\left(s\right)$ helps to render the structure s with spatially sparse features by adding a penalty based on the absolute value of s as follows,

$${ℛ}_{\mathrm{L}1}\left(s\right)=\mathrm{l}\mathrm{o}\mathrm{g}\left(\left|s\right|+{\epsilon }_{L1}\right)\simeq \frac{\left|s\right|}{\mathrm{s}\mathrm{g}\left(\left|s\right|\right)+{\epsilon }_{L1}}.$$ (14)

Here, L1 regularization is applied to an entire 3D volume. ${\epsilon }_{L1}={\sigma }_b\left(\left|g_{bfr}\right|\right)$ and the same logarithmic tone mapping function²⁵ (Eq. 12) is applied on the top of the absolute value.

Lastly, NLD regularization³⁹ ${ℛ}_{\mathrm{N}\mathrm{L}\mathrm{D}}\left(s\right)$ constrains the first-order difference of the structure s along the depth axis z. It prevents axial features of the structure from fitting to the rapidly varying spatial components that sparsity-promoting regularizations (Equations 11, 14) might favor. This regularization balances the influence of the first two terms, allowing the structure to retain desirable details along the z-axis. It is written as

$${ℛ}_{\mathrm{N}\mathrm{L}\mathrm{D}}\left(s\right)={\left.\left|\frac{\partial s}{\partial z}\right|\right|}_{\delta ,\left[a,b\right]},$$ (15)

where ${\left.f\right|}_{\delta ,\left[a,b\right]}\equiv \mathrm{m}\mathrm{a}\mathrm{x}\left(f,b\right)+\delta \mathrm{m}\mathrm{a}\mathrm{x}\left(a,\min \left(f,b\right)\right)+\mathrm{m}\mathrm{i}\mathrm{n}(f,a)$. For all results presented in this manuscript, $\delta =0.1$, a = 0.005, b = 2.0.

Together, the summation of the regularization terms is expressed as

$$ℛ\left(s\right)={\lambda }_{tv}{ℛ}_{tv}\left(s\right)+{\lambda }_{\mathrm{L}1}{ℛ}_{\mathrm{L}1}\left(s\right)+{\lambda }_{\mathrm{N}\mathrm{L}\mathrm{D}}{ℛ}_{\mathrm{N}\mathrm{L}\mathrm{D}}\left(s\right),$$ (16)

where ${\lambda }_{tv}=0.005$, ${\lambda }_{\mathrm{L}1}=0.01$, ${\lambda }_{\mathrm{N}\mathrm{L}\mathrm{D}}={10}^{-6}$.

Baseline term b in image formation

The baseline term b is modeled as a low-rank component to capture background fluorescence, which primarily arises from tissue autofluorescence and is typically slowly varying across the image. b is represented as the sum of rank-1 tensors, and we set to R = 5 here:

$$b=\sum _{r=1}^R{b}_{z,r}\times b_{y,r}\times b_{x,r},$$ (17)

where $b_{x,r}$, $b_{y,r}$, $b_{z,r}$ are learnable 2D tensors to represent baseline components along the x, y, and z axes, respectively. These tensors are initialized with the value ${\left(0.1{\sigma }_b\left(\left|g_{bfr}\right|\right)\right)}^{1/3}$. By constraining b to low rank, we limit it to low-spatial-frequency features, effectively separating background fluorescence from fluorescent features of interest. For input stack g acquired from a mouse brain slice with GFP-expressing neurons (Fig. S14a) and in vivo from a mouse brain with iGluSnFR3-expressing neurons (Fig. S14b), respectively, the fluorescence baseline is much dimmer than labeled neurons and its low-rank nature is obvious. For the brain slice sample, autofluorescence only exists within the tissue slice and decrease close to zero outside the tissue (left edge of the yz image for Baseline plotted on a reduced scale in Fig. S14a).

Two-step learning process

The weights of the neural network $\theta$ (representing structure s), Zernike coefficients $\alpha$ (thus PSF $h\left(\mathit{r};\alpha \right)$), and baseline term b in the image-formation model are optimized in a two-step learning process³⁵. The first step only adjusts neural network weights for s, while the Zernike coefficients $\alpha$ and baseline b remain fixed after initialization, where $\alpha$ is randomly initialized with a random value from a uniform distribution in the range [0, 10⁻²]. It conditions the randomly initialized neural network, using the loss function:

$${\theta }^{\text{*}}=\underset{\theta }{\text{argmin}}\left(1-\text{SSIM}\left(c{g}_{lp},f_{\theta }\left(r\right)\right)\right),c\ge 1$$ (18)

where $g_{lp}$ is a low pass filtered image stack with an isotropic Gaussian filter. Optimization is performed using the RAdam optimizer⁶⁰ with an initial learning rate of 10⁻², ${\beta }_1=0.9$, and ${\beta }_2=0.999$ for 5000 epochs. The learning rate schedule follows an exponential decay down to 10⁻³ by the end of the epoch.

The second step updates neural network weights $\theta$, Zernike coefficients $\alpha$, and baseline b using the loss function (Eq. 3). For this learning process, the initial learning rate is set to 4×10⁻³ with the same RAdam optimizer, keeping ${\beta }_1$ and ${\beta }_2$ unchanged, running for 5000 epochs. The learning rate schedule again follows an exponential decay, this time down to 10⁻⁶ by the end of the epoch.

All computational implementations are performed on a machine equipped with an NVIDIA RTX 4090 GPU, an Intel i9-13900K CPU, and 80 GB of RAM. The computation time for the results in the main figures is listed in Table S2, along with their corresponding experimental settings. The scaling of NeAT’s computational time with respect to the input image size is described in Supplementary Note 2 and Figs. S15, S16. Overall, the computational cost of NeAT scales approximately linearly with the square root of total number of pixels in the XY images of the input stack, while remaining constant with respect to the axial depth (for a fixed number of z-slices).

Preprocessing of 3D image stacks from in vivo experiments

The raw 3D experimental fluorescence image stacks have dimensions of $N_f\times N_z\times N_y\times N_x$, where N_f denotes the number of frames per z-axis slice (typical N_f ~ 50), $N_z$ the number of z-axis slices, and $N_x$ and $N_y$ the number of pixels along the x- and y-axes, respectively. Here $N_f$ frames are acquired per z-axis slice to reduce the effect of Gaussian noise through averaging. In a typical experiment, N_f = 50 frames were acquired at a z-axis slice, before advancing to the next z-slice and acquiring another N_f frames. Collecting N_z (typically 50) z-axis slices in this manner required approximately 1.5 minutes in total. In in vivo imaging experiments, the frames acquired at the same z may need to be registered before averaging to correct for sample motion between frames. We used a customized ImageJ plugin (Code S1 in Supplementary Material) to register the frames for each z-axis slice in 2D using the TurboReg plugin with a rigid body assumption and then to average the N_f registered frames to obtain the image stack with dimensions of $N_z\times N_y\times N_x$. The resulting stack is cropped to remove edge pixels and then used as input to NeAT. The input stack dimensions listed in Table S2 of the Supplementary Material reflect this preprocessing step. For the in vivo image stacks shown in Figs. 4 and 5, the dimensions are set to 50 × 200 × 200.

Typical input image stack extends 10 μm in z, within which aberrations are effectively constant, an assumption supported by observations that applying an aberration correction estimated at one depth substantially improves signal intensity and contrast across adjacent depths spanning at least ~100 μm (i.e., ±50 μm)^{43,48,49,61,62}.

We explored whether denoising by Noise2Void⁶³ could reduce frame averaging requirements, but found that at the SNR of our in vivo images, such a computational denoising approach caused errors in aberration estimation when the number of averaged frames was too low (Fig. S17, Supplementary Note 3). Therefore, for the SNR ranges that we explored, frame averaging remains the preferred approach for noise reduction.

Motion correction of input image stacks by NeAT

Image stacks acquired in vivo can contain motion artifacts caused by heartbeats, breathing, or body movements. Although preprocessing as described above removes the motion artifacts for frames acquired at the same z depth, motion between frames at different z depths also needs to be corrected to ensure that structural features are properly aligned for accurate aberration estimation by NeAT. Motion correction for the input z-stack therefore refers to registration across different z positions. Failing to so would lead to errors in aberration estimation and structural recovery (Fig. S5). NeAT incorporates motion correction across z slices into its learning process by assigning an affine transformation matrix, $A_{n_z}\left(n_z=1,2,\cdots ,N_z\right)$, to each z slice to correct translation, rotation, scaling, and shear caused by the sample’s motion. This is formulated as $\left(Ag\right)\left[n_z\right]=A_{n_z}g\left[n_z\right]$, where $g\left[n_z\right]$ denotes the n_z-th z-slice of the input image stack g. NeAT corrects motion by iteratively updating these matrices during the learning process. Each matrix is a 2×3 matrices with learnable elements and initialized as the identity. We used the RAdam optimizer⁶⁰ for the motion correction process, with an initial learning rate of 0.07, ${\beta }_1=0.9$, and ${\beta }_2=0.999$. More details are available in our public repository. Through aberration-aware motion registration, for image stacks with aberration, NeAT leads to different affine parameters from those estimated by StackReg in ImageJ (Fig. S18) and achieves more accurate aberration measurement than pre-registering the image stack prior to learning using StackReg (Fig. S11).

Calculation of signal-to-noise ratio

We assumed a linear relationship between the grayscale pixel value (d) and the photon count per pixel (p_c), expressed as $d=\beta p_c$, where β is the conversion factor. To compute β, we acquired multiple images (e.g., more than 100) of a fluorescein solution at the same imaging condition. We then calculated the variance and mean for d. Since p_c theoretically follows a Poissonian distribution, where the variance equals the mean, $\beta$ is computed as the ratio of the variance to the mean of d.

For the cutoff SNR analysis, we calculated $\beta$ for the PMT in the custom-build microscope under different control voltages, observing gains of 7.83 at a control voltage of 0.7 V (used for acquiring images from 1-μm fluorescence beads, Figs. 2j–l) and 21.8 at a control voltage of 0.8 V (used for fixed Thy1-GFP mouse brain slice imaging, Figs. 2m–o).

Next, we classified the pixels in an image stack as either signal or background pixels using a classification method described previously³⁵. We then calculated the SNR of the image stack as

$$\mathrm{S}\mathrm{N}\mathrm{R}=\frac{\stackrel{-}{y}/\beta }{\sqrt{\stackrel{-}{y}/\beta }}=\sqrt{\stackrel{-}{y}/\beta },$$ (19)

where $\stackrel{-}{y}$ is the mean grayscale value of the signal pixels, and $\stackrel{-}{y}/\beta$ represents the corresponding photon count. We developed an ImageJ plugin (Compute_SNR_2P.py; Code S2 in Supplementary Material) to compute the SNR of a 3D image stack and determine whether it possesses sufficient SNR to be used as NeAT input (i.e., its SNR should exceed the SNR cutoff).

In vivo imaging of visually evoked glutamate and calcium activity

Visual stimuli were generated in MATLAB using the Psychophysics Toolbox^64,65 and presented 15 cm from the left eye of the mouse on a gamma-corrected, LED-backlit LCD monitor with a mean luminance of 20 cd·m⁻². We divided the monitor into a 3×3 grid and presented 1-s-long uniform flashes in a pseudorandom sequence in one of the 9 grids, while recording fluorescence images with a 2 mm by 2 mm FOV. Analyzing these images allowed us to identify the cortical region that responded to the center of the monitor. We then imaged this cortical region at smaller pixel sizes to measure glutamate and calcium activity of synapses and neurons towards oriented drifting grating stimulation in mice under light anesthesia (0.5% isoflurane in O₂). Full-field gratings of 100% contrast, a spatial frequency of 0.04 cycles per degree, and a temporal frequency of 2 Hz drifting in eight directions (0° to 315° at 45° increments) were presented in pseudorandom sequences. For glutamate imaging (x and y pixel size: 0.125 μm/pixel), each grating stimulus lasted 2 s with a 1-s presentation of a gray screen before and after the stimulus. For calcium imaging (x and y pixel size: 0.945 μm/pixel), each grating stimulus lasted 2 s with a 1-s gray screen presentation before and a 3-s gray screen presentation after the stimulus. Each stimulus was repeated for 10 trials per imaging session.

Functional image analysis

Images were processed using custom Python code (see Code S3 in the Supplementary Material.) Glutamate time-lapse images were registered using iterative phase correlation with polar transform (implemented using the scikit-image Python package) to correct for non-rigid motions including translation, rotation, and scaling. The images from the first trial of visual stimulation (a 4-second-long recording) were used as the reference. Within this reference time series, each frame was registered to the first frame. Then an average was calculated from all registered frames therein and used as the reference. Images from subsequent trials were registered to this reference. Calcium time-lapse images were registered with the StackReg package⁶⁶, with the first frame as the reference. In our case, the polar transform method proved more effective than StackReg for registering glutamate images, which were dimmer and noisier than calcium images. Regions of interest (ROIs) were manually drawn in ImageJ using the circular selection tool on the mean intensity projection of the glutamate time-lapse images and elliptical selection tool for the GCaMP6s time-lapse images. The ROIs were then imported into a Python environment to extract pixel values within the ROIs, which were averaged to obtain the raw fluorescence signal F for each ROI.

The glutamate transient $∆F/F_0$ was calculated as $(F-F_0)/F_0$, where $F_0$ represents the basal fluorescence, defined as the average fluorescence signal during the 1-s pre-stimulus gray-screen presentation period, excluding the highest 5% of values in F from the calculation.

For calcium images, due to higher labeling density, we removed neuropil contamination. We calculated $F_{\text{neuropil}}$ as the averaged fluorescence signal from the neuropil area⁴⁷ (defined as the pixels that were 2 to 20 pixels off the ROI border) and computed $∆F_{\text{neuropil}}$ as $F_{\text{neuropil}}-F_{0,\text{neuropil}}$, where $F_{0,\text{neuropil}}$ is the mean of $F_{\text{neuropil}}$ during the 1-s pre-stimulus period. Then, $∆F_{\text{neuropil}}$ was multiplied by 0.7 and subtracted from F to obtain $F_{\text{true}}$. $∆F/F_0$ was then computed as $(F_{\text{true}}-F_{0,\text{true}})/F_{0,\text{true}}$, with $F_{0,\text{true}}$ defined as the mean of $F_{\text{true}}$ during the 1-s pre-stimulus period.

Trial-averaged $∆F/F_0$ was calculated as the average of 10 trials. Peak $∆F/F_0$ was defined as the maximal trial-averaged $∆F/F_0$ within the 2-s drifting grating presentation. Response R for each drifting grating direction was defined as the averaged $∆F/F_0$ across the 2-s drifting-grating stimulus presentation, with negative responses set to zero.

For glutamate images, an ROI was considered responsive to visual stimulation if its peak $∆F/F_0$ was greater than 3 times the standard deviation of the trial-averaged $∆F/F_0$ within the 2-s stimulus period^67,68 and if the peak $∆F/F_0$ was above 5%⁶⁹. For calcium images, an ROI was considered active if its maximal $∆F/F_0$ was above 10%^57,70 and visually responsive if its activity during at least one visual stimulus type was significantly higher than its activity during the pre-stimulus period, as determined by one-way ANOVA with $p<0.01.$ All traces shown in Fig. 5 were filtered using a Savitzky–Golay filter⁴⁴.

Orientation selectivity analysis

For each ROI, its tuning curve $R_{fit}\left(\theta \right)$ was defined as the fitted curve to $R\left(\theta \right)$ with a bimodal Gaussian function⁴⁵:

$$R_{fit}\left(\theta \right)=R_0+A_1{e}^{-\frac{ang{\left(\theta -{\theta }_{pref}\right)}^2}{2{\sigma }^2}}+A_2{e}^{-\frac{ang{\left(\theta -{\theta }_{pref}+{180}^{\circ }\right)}^2}{2{\sigma }^2}},$$ (20)

where $ang\left(x\right)=\mathrm{m}\mathrm{i}\mathrm{n}\left(\left|x\right|,\left|x-{360}^{\circ }\right|,\left|x+{360}^{\circ }\right|\right)$, which wraps the angular values onto the interval between $0^{\circ }$ and ${180}^{\circ }$. Responses to the different drifting direction $R\left(\theta \right)$ were fitted to the function to minimize the mean square error between the model and responses, with $R_0$, $A_1$, $A_2$ constrained to non-negative values, and $\sigma$ constrained to be larger than 22.5 ⁷¹, given that the angle step was 45.

ROIs were considered orientation-sensitive (OS) if their responses across 8 different drifting grating stimuli were significantly different by one-way ANOVA ($p<0.05$)^46,69 and if their responses were well-fit to the bimodal Gaussian model⁷⁰. The goodness of the fit was assessed by calculating the error $E$ and the coefficient of determination ${\mathfrak{R}}^2$:

$$E=\sum _{n=0}^7{{\left(R\left(\theta \right)-R_{fit}\left(\theta \right)\right)}^2|}_{\theta ={\left(45n\right)}^{\circ }},{\Re }^2=1-\frac{E}{{\sum }_{n=0}^7{{\left(R\left(\theta \right)-\overline{R}\right)}^2|}_{\theta ={\left(45n\right)}^{\circ }}},$$ (21)

where $\stackrel{-}{R}$ is the mean of $R\left(\theta \right)$. The criteria for a good fit were $E<0.01$ and ${\mathfrak{R}}^2>0.5$. The fitted response was used to calculate orientation sensitivity index (OSI) as $\frac{R_{\text{pref}}-R_{\text{ortho}}}{R_{\text{pref}}+R_{\text{ortho}}}$, where $R_{\text{pref}}$ and $R_{\text{ortho}}$ are the responses at ${\mathrm{\theta }}_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{f}}$ and ${\theta }_{\text{ortho}}\left(={\mathrm{\theta }}_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{f}}+{90}^{\circ }\right)$, respectively.

Statistics

Standard functions from the Scipy package in Python were used to perform statistical tests, including two-sided paired t-test, one-way ANOVA, and Kolmogorov-Smirnov test. Statistical significance was defined as *$p<0.05$, ***$p<0.01$, and ***$p<0.001$.

Data availability

Data for conjugation error estimation/correction and aberration estimation from both brain-slice and in vivo experiments are available on Code Ocean under DOI: 10.24433/CO.2804276.v1.

Code availability

Code is available on Code Ocean under the DOI: 10.24433/CO.2804276.v1, along with metadata describing its functionality.