Automated discovery of experimental designs in super-resolution microscopy with XLuminA

Abstract

Driven by human ingenuity and creativity, the discovery of super-resolution techniques, which circumvent the classical diffraction limit of light, represent a leap in optical microscopy. However, the vast space encompassing all possible experimental configurations suggests that some powerful concepts and techniques might have not been discovered yet, and might never be with a human-driven direct design approach. Thus, AI-based exploration techniques could provide enormous benefit, by exploring this space in a fast, unbiased way. We introduce XLuminA, an open-source computational framework developed using JAX, a high-performance computing library in Python. XLuminA offers enhanced computational speed enabled by JAX's accelerated linear algebra compiler (XLA), just-in-time compilation, and its seamlessly integrated automatic vectorization, automatic differentiation capabilities and GPU compatibility. XLuminA demonstrates a speed-up of 4 orders of magnitude compared to well-established numerical optimization methods. We showcase XLuminA's potential by re-discovering three foundational experiments in advanced microscopy, and identifying an unseen experimental blueprint featuring sub-diffraction imaging capabilities. This work constitutes an important step in AI-driven scientific discovery of new concepts in optics and advanced microscopy.

Fig. 1. Overview and performance of XLuminA.

a Software’s workflow shows integrated feedback between the AI discovery tool and the optics simulator. Initial random optical parameters shape the hardware design on a virtual optical table. The optics simulator computes the performance of the experiment through detected light, from which the objective function (for instance, the spot size ϕ, where FWHM stands for Full Width Half Maximum) is evaluated. To improve the cost function metric, the optimizer adjusts the optical parameters, creating an iterative cycle between the simulator and optimizer until convergence. b Average execution time (in seconds) over 100 runs at 2048 × 2048 pixel resolution, for scalar and vectorial field propagation using Rayleigh-Sommerfeld (RS, VRS) and Chirped z-transform (CZT, VCZT) algorithms in Diffractio and XLuminA. Using pre-compiled jitted functions, XLuminA achieves × 2 speedup for RS and CZT and  × 2.5 for VRS and VCZT on CPU. GPU implementation improves the performance up to × 64 for RS, × 76 for CZT, × 80 for VRS, and × 78 for VCZT. c Average time (in seconds) over 5 runs for a single gradient evaluation using numerical differentiation (num. diff) with Diffractio’s optical simulator (blue dots) and auto-differentiation (autodiff) methods (green triangles for CPU and magenta squares for GPU) with XLuminA's optical simulator for different resolutions. At 250 × 250 pixel resolution, GPU-based XLuminA’s autodiff methods significantly outperform numerical methods by a factor of × 3.9 ⋅ 10⁵, and a factor of × 1.8 ⋅ 10⁴ in the CPU. d Average time (in seconds) over 5 runs for convergence time, using numerical differentiation with Diffractio’s optical simulator and autodiff methods with XLuminA’s optical simulator for different resolutions. At 250 × 250 pixel resolution, GPU-based XLuminA’s autodiff methods significantly outperform numerical methods by a factor of × 2.1 ⋅ 10⁴ and a factor of × 8.4 ⋅ 10² in the CPU. Standard deviation corresponds to shaded regions. We use BFGS and Adam optimizers, for numerical and autodiff approaches, respectively. The superior efficiency of autodiff over traditional numerical methods allows for highly efficient optimizations, particularly employing the large high resolutions we use (up to 2048 × 2048 pixels).

Fig. 2. Rediscovery of the optical configuration employed to magnify images.

a Virtual optical arrangement. It consists of a light source emitting a 650 nm wavelength Gaussian beam. Original lenses are replaced by two spatial light modulators (SLMs) with a resolution of 1024 × 1024 and a pixel size of 2.92 μm. The parameter space (of ∼2 million optical parameters) includes the distances, z₁, z₂, and z₃ (in millimeters) and the phase masks (in radians) of the two SLMs. b Data-driven discovery scheme. Input-output sample pairs are fed into the optics simulator in batches of 10. The loss function, computed for each virtual optical setup, evaluates the mean squared error between the intensity response of the system and the corresponding target example from the dataset. The average loss over the batch guides the optical parameter update, which is common to all the virtual optical setups. This cycle is repeated until convergence is reached. c Identified phase mask solutions for SLM#1 and SLM#2. Identified distances correspond to z₁ = 10.14 cm, z₂ = 5.46 cm and z₃ = 7.54 cm. Input, detected intensity, and expected (ground truth) intensity patterns for (d) a simple geometry, and (e) a complex structure, the Max Planck Society’s logo. In both cases, the identified optical design successfully inverts and magnifies 2 × the input mask.

Fig. 3. General virtual optical setup for large-scale discovery schemes.

Gray boxes represent fundamental building units, each containing a super-SLM and a wave plate positioned a distance z₁ apart. These units are inter-connected through free propagation distances z₂, and beam splitters (PBS). The super-SLM is a hardware box type that consists of two spatial light modulators (SLMs), each one independently imprinting a phase pattern on the horizontal and vertical polarization components of the field. The setup’s complexity and size can be arbitrarily extended by incorporating additional connections, building units, light sources, detectors, etc. The quantization of the large-scale search space generated by this optical setup is provided in the Methods section.

Fig. 4. Pure topological discovery within a fully continuous framework.

The parameter space (25 optical parameters) is defined by 9 beam splitter ratios (BS), 8 distances (z), and 4 wave plates (WP, with variable phase retardance η and orientation angle θ). a Discovered optical topology for STED microscopy. The minimum value of the loss function is demonstrated in detector #3. The phase mask (PM) #2 corresponds to the radial phase pattern originally used in STED microscopy, which generates a doughnut-shaped beam. b Radial intensity profile, ∣E_x∣² + ∣E_y∣², in horizontal beam section: excitation (green), depletion (orange), and super-resolution effective STED beam (dashed blue line). Lateral position indicates lateral distance from the optical axis. The data corresponding to the original STED phase are indicated with dotted lines. The excitation and depletion beams are diffraction-limited. The effective response breaks the diffraction limit. c Discovered optical topology for Dorn, Quabis, and Leuchs (2003). The minimum value of the loss function is demonstrated in detector #3. Phase masks #1 and #2 correspond to the polarization converter demonstrated in Dorn, Quabis, and Leuchs (2003) and a radial phase pattern originally used in STED microscopy, respectively. Both phase patterns generate, independently, a doughnut-shaped beam. d Normalized longitudinal intensity profile, ∣E_z∣², for Dorn, Quabis, and Leuchs (2003) and the identified solution (black dotted, and green lines, respectively) and radial intensity profile, ∣E_x∣² + ∣E_y∣², of the diffraction-limited linearly polarized beam (orange dotted line). Lateral position indicates lateral distance from the optical axis. The spot size is computed as ϕ = (π/4)FWHM_x FWHM_y, where FWHM denotes Full Width Half Maximum. The discovered approach breaks the diffraction limit, demonstrating a spot size close to the reference. Also, it does not feature side lobes (indicated with a gray arrow), which can limit practical imaging techniques. Details of the initial optical setup, phase masks, and identified optical parameters can be found in Supplementary Information section Rediscovery Through Exploration of Optical Topologies.

Fig. 5. Rediscovery of STED microscopy within highly parameterized optical systems.

a Discovered optical topology. The parameter space  (∼4 million optical parameters) is defined by 3 super-SLMs (i.e., 6 SLMs) of (824 × 824) pixel resolution and a computational pixel size of 6.06 μm, 9 beam splitter ratios (BS), 8 distances (z) and 3 wave plates (with variable phase retardance η and orientation angle θ). The minimum value of the loss function is demonstrated in detector #2. The setup topology is retrieved from detector #2 following the identified beam splitter ratios across the system. The identified optical parameters correspond to: the beam splitter ratios, in [Transmittance, Reflectance] pairs: BS#1: [0.000, 0.999], BS#2: [0.201, 0.799], BS#5: [0.000, 0.999], and BS#6: [0.999, 0.000]. The wave plates, in radians (1): η = − 1.39, θ = − 1.64, and (2): η = − 1.61, θ = − 0.86. The propagation distances (in cm) are z₁ = 59.52, z₂= 10.14, z₃ = 76.36, z₄ = 17.93, z₅ = 37.07, z₆ = 65.95, and z₇ = 38.68. b Discovered phase patterns for sSLM #1 and sSLM #2. The speckle-like patterns of SLMs' phase masks are not detrimental to robustness. c Radial intensity profile, ∣E_x∣² + ∣E_y∣², in horizontal beam section: excitation (green), depletion (orange), and super-resolution effective STED beam (dashed blue line). The data corresponding to the original STED experiment - i.e., computed using a spiral phase mask - are indicated with dotted lines. Lateral position indicates lateral distance from the optical axis.

Fig. 6. Rediscovery of Dorn, Quabis, and Leuchs (2003) within highly parameterized optical systems.

a Discovered virtual optical setup topology. The parameter space  (∼6.2 million optical parameters) is defined by 3 super-SLMs (i.e., 6 SLMs) of (1024 × 1024) pixel resolution and a computational pixel size of 4.8 μm, 9 beam splitter ratios (BS), 8 distances (z) and 3 wave plates (with variable phase retardance η and orientation angle θ). The minimum value of the loss function is demonstrated in detector #6. The setup topology is retrieved from detector #6 following the identified beam splitter ratios across the system. The identified optical parameters correspond to: the beam splitter ratios, in [Transmittance, Reflectance] pairs: BS#2: [0.999, 0.000], BS#5: [0.000, 0.999], BS#6: [0.000, 0.999], and BS#9: [0.999, 0.000]. The wave plate’s η = 1.51, θ = 3.95; propagation distances (in cm): z₃ = 20.49, z₄ = 63.26, z₇ = 47.92 and z₈ = 31.33. b Discovered phase patterns for sSLM #2. The speckle-like patterns of SLMs' phase masks are not detrimental to robustness. c Normalized longitudinal intensity profile, ∣E_z∣², for Dorn, Quabis, and Leuchs (2003) and the identified solution (black dotted, and green lines, respectively) and radial intensity profile, ∣E_x∣² + ∣E_y∣², of the diffraction-limited linearly polarized beam (orange dotted line). Lateral position indicates lateral distance from the optical axis. The spot size is computed as ϕ = (π/4)FWHM_x FWHM_y, where FWHM denotes Full Width Half Maximum. The discovered approach breaks the diffraction limit, demonstrating a larger spot size as the reference. However, it does not feature side lobes (indicated with a gray arrow), which can limit practical imaging techniques.

Fig. 7. Discovery of a previously unreported experimental blueprint within a highly parameterized optical setup.

The parameter space comprises 6 × (824 × 824) pixel phases, 6 extra optical modulation parameters corresponding to the wave plates, and 8 distances (a total of ∼4 million optical parameters). a Discovered optical topology. The minimum value of the loss is demonstrated in detector #2. The setup topology is easily retrieved from detector #2 following the identified beam splitter ratios across the system. The identified optical parameters correspond to: the beam splitter ratios, in [Transmittance, Reflectance] pairs: BS#1: [0.000, 0.999], BS#2: [0.338, 0.662], BS#5: [0.000, 0.999], and BS#6: [0.999, 0.000]. The wave plates, in radians (1): η = 1.09, θ = 0.28, and (2): η = 0.19, θ = − 3.16. The propagation distances (in cm) are z₁ = 29.98, z₂ = 56.91, z₃ = 58.28, z₄ = 99.91, z₅ = 42.89, z₆= 53.96, and z₇ = 50.96. b Discovered phase masks corresponding to the super-SLM (sSLM) in (1) and (2). The speckle-like patterns of SLMs' phase masks are not detrimental to robustness. c Total intensity (∣E_x∣² + ∣E_y∣² + ∣E_z∣²) horizontal cross-section of the detected light beams of 650 nm (orange), 532 nm (green), and effective beam emulating stimulated emission (dashed blue). d Horizontal cross-section of the normalized total intensity (∣E_x∣² + ∣E_y∣² + ∣E_z∣²) of the effective beam from the discovered solution (blue), the simulated STED reference (dashed red), and the simulated reference (dotted black) using 532 nm wavelength. The spot size is computed as ϕ = (π/4)FWHM_x FWHM_y, where FWHM denotes Full Width Half Maximum. We use a fluorophore emission wavelength of 560 nm. The discovered solution outperforms both simulated references for STED microscopy and the sharp focus from Dorn, Quabis, and Leuchs (2003), demonstrating a spot size of 9.45% smaller than both references.

Fig. 8. Performance of the discovered optical configuration employed to magnify the image under different levels of noise and misalignment.

a Input mask intensity. b Ground truth intensity pattern (2 × magnification). c Detected intensity pattern using the discovered optical system with ideal optical components and perfect alignment. d Performance under standard experimental conditions: uniformly distributed random noise of ± (0.01 to 0.1) radians in spatial light modulators (SLMs), and ± (0.01 to 0.1) mm (millimeters) of misalignment. e Performance under very large uniformly distributed random noise of ± (0.1 to 1) radians in SLMs and ± (0.1 to 1) mm of misalignment. f Performance under very large uniformly distributed random noise of ± (0.1 to 1) radians, discretized at 8-bit, in SLMs and ± (0.1 to 1) mm of misalignment. The discovered system maintains good imaging performance even under very high error conditions, demonstrating its robustness.