https://doi.org/10.5061/dryad.wdbrv15wk

Python codes for performing wavefront optimization with the sparsity constraints associated with the dimensionality limitation of spatial light modulation devices.

Overview

Wavefront shaping gains increasing importance in complex photonics, which can manipulate light spatially and temporally to counter the scattering effect. Important applications include deep-tissue imaging, microendoscopy, optical communications, nanofabrication, and remote sensing. However, high-speed and high-fidelity wavefront shaping is fundamentally hindered by the dimensionality limitation of hardware devices, evinced by the competition between the frame rate, pixel count, and modulation depth. To overcome the speed-fidelity tradeoff, we leverage complex media (e.g., diffusers or multimode fibers) as analogue random multiplexers for pattern compression to address the demand for high-dimensional spatiotemporal control. Sparsity-constrained wavefront optimization is designed to solve the problem by seeking a low-dimensional, robust representation of wavefronts with carefully designed sparsity constraints. This optimization framework can achieve high-fidelity wavefront shaping through complex media using high-speed, yet relatively low-precision spatial light modulation devices (e.g., digital micromirror devices) without compromising the frame rate.

Installation

The codes require the installation of CuPy and PyTorch to execute. The codes have been tested and verified with Python 3.9.

Descriptions

This dataset contains the preprocessed test images (data/fashion_dataset_1000.hdf5), an experimentally calibrated transmission matrix (data/projection_optimization/complexTM.hdf5), an auxiliary file for specifying the output coordinates (data/projection_optimization/cam_roi.npz), and two scripts for benchmarking the performance of the vanilla PyTorch implementation and the customized CuPy implementation (gradient_descent_cupy.py and gd_implementation_benchmark.py)

data/fashion_dataset_1000.hdf5

• xs_init: initial guess of wavefronts (obatined by Gerchberg-Saxton algorithm)
• yts: preprocessed target output patterns extracted from Fashion-MNIST dataset.
• ss: scaling factor for normalization in the mean square error calculation.

data/projection_optimization/complexTM.hdf5

• T1, T2: transmission matrices of two orthogonal linear polarization states at the distal end.

data/projection_optimization/cam_roi.npz

• pts1_phy, pts2_phy: the output coordinates of the two polarization states in unit of um. In general, two coordinates are identical.
• r: the radius of the output speckles in unit of camera pixel.

gradient_descent_cupy.py

• adam(x, forward_func, loss_func, grad_func, num_iters, alpha, beta1=0.9, beta2=0.999, eps=1e-8, **kwargs): customized ADAM optimizer.

gd_implementation_benchmark.py

• forward_func(T): CuPy-implemented forward model of light propagation through a complex medium given a transmission matrix T.
• mseloss(y, yt, s): Cupy-implemented mean squared error between the estimated output y and the target yt, with a learnable scalar s.
• loss_func(x, yc, **kwargs): Cupy-implemented loss function of the sparsity-constrained wavefront optimization.
• gradient_comp(x, y, **kwargs): customized gradient calculation for the sparsity-constrained wavefront optimization.
• forward_func_torch(T): PyTorch-implemented forward model.
• gradient_descent_pytorch(u_est, y_tar, num_iters, forward, loss=nn.MSELoss(), lr=0.01, reg=0.01, s0=1.0, device='cpu'): gradient descent method with vanilla PyTorch implementation.

Example

In gd_implementation_benchmark.py, the computation speed of the customized $l_1$-regularized gradient descent method based on CuPy is benchmarked in comparison with the vanilla PyTorch implementation. The customized CuPy implementation is included in function adam in gradient_descent_cupy.py. An experimentally calibrated transmission matrix of a graded-index multimode fiber (GIF50C, Thorlabs) and the preprocessed target patterns extracted from the Fashion-MNIST dataset are loaded, followed by the wavefront optimization using two different implementations. The runtime and the optimization curve are displayed in the end for the performance comparison.

Reference

L. Yu, and S. You, "High-Fidelity and High-Speed Wavefront Shaping by Leveraging Complex Media," Science Advances, in press. Preprint available at arXiv:2302.10254. https://doi.org/10.48550/arXiv.2302.10254

Sharing/Access permission

MIT License
Copyright (c) 2023 Li-Yu Yu
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

Disclaimer:

The software is provided "as is", without warranty of any kind, express or implied, including but not limited to the warranties of merchantability, fitness for a particular purpose and noninfringement. In no event shall the authors or copyright holders be liable for any claim, damages or other liability, whether in an action of contract, tort or otherwise, arising from, out of or in connection with the software or the use or other dealings in the software.