Package: xlumina
✨ XLuminA ✨XLuminA, a highly-efficient, auto-differentiating discovery framework for super-resolution microscopyXLuminA: An Auto-differentiating Discovery Framework for Super-Resolution MicroscopyCarla Rodríguez, Sören Arlt, Leonhard Möckl and Mario Krenn💻 Installation:XLuminA can be installed withpip install xluminaThis will installJAX and jaxlib(the version of JAX used in this project is v0.4.13),Optax(the version of Optax used in this project is v0.1.7), andSciPy(the version of SciPy used in this project is v1.10.1).GPU compatibility:The package automatically installs the CPU version of JAX. To installJAX with NVIDIA GPU support(Note: wheels only available on linux), use:# CUDA 12 installation
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11 installation
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html👾 Features:XLuminA allows for the simulation, in a (very) fast and efficient way, of classical light propagation through optics hardware configurations,and enables the optimization and automated discovery of new setup designs.The simulator contains many features:✦ Light sources (of any wavelength and power) using both scalar or vectorial optical fields.✦ Phase masks (e.g., spatial light modulators (SLMs), polarizers and general variable retarders (LCDs)).✦ Amplitude masks (e.g., circles, triangles and squares).✦ Beam splitters.✦ The light propagation methods available in XLuminA are:Fast-Fourier-transform (FFT) based numerical integration of the Rayleigh-Sommerfeld diffraction integral.Chirped z-transform. This algorithm is an accelerated version of the Rayleigh-Sommerfeld method, which allows for arbitrary selection and sampling of the region of interest.Propagation throughhigh NA objective lensesis availale to replicate strong focusing conditions in polarized light.📝 Example of usage:Examples of some experiments that can be reproduced with XLuminA are:Optical telescope (or 4f-correlator),Polarization-based beam shaping as used inSTED (stimulated emission depletion) microscopy,Thesharp focus of a radially polarized light beam.The code for each of these optical setups is provided in the Jupyter notebook ofexamples.ipynb.🚀 Testing XLuminA's efficiency:We evaluated our framework by conducting several tests - seeFigure 1. The experiments were run on an Intel CPU Xeon Gold 6130 and Nvidia GPU Quadro RTX 6000.(1) Average execution time (in seconds) over 100 runs, within a computational window size of $2048\times 2048$, for scalar and vectorial field propagation using Rayleigh-Sommerfeld (RS, VRS) and Chirped z-transform (CZT, VCZT) inDiffractioand XLuminA. Times for XLuminA reflect the run with pre-compiled jitted functions. The Python files corresponding to light propagation algorithms testing arescalar_diffractio.pyandvectorial_diffractio.pyfor Diffractio, andscalar_xlumina.pyandvectorial_xlumina.pyfor XLuminA.(2) we compare the gradient evaluation times of numerical (using SciPy'sBFGS optimizer)vsanalytical differentiation (using JAX'sADAM optimizer) when optimizing using XLuminA's optical simulator...(3) and its convergence time:The Jupyter notebook used for running these simulations is provided astest_diffractio_vs_xlumina.ipynb.The Python files corresponding to numerical/autodiff evaluations arenumerical_methods_evaluation_diffractio.py,numerical_methods_evaluation_xlumina.pyandautodiff_evaluation_xlumina.pyIf you want to run the comparison test of the propagation functions, you need to installDiffractio- The version of Diffractio used in this project is v0.1.1.🤖🔎 Discovery of new optical setups:With XLuminA we were able to re-discover three foundational optics experiments:➤ Optical telescope (or 4f-correlator),➤ Polarization-based beam shaping as used inSTED (stimulated emission depletion) microscopy,➤ Thesharp focus of a radially polarized light beam.The Python files used for the discovery of these optical setups, as detailed inour paper, are organized in pairs ofoptical_tableandoptimizeras follows:Experiment name🔬 Optical table🤖 Optimizer📄 File for dataOptical telescopefour_f_optical_table.pyfour_f_optimizer.pyGenerate_synthetic_data.pyPolarization-based STEDvsted_optical_table.pyvsted_optimizer.pyN/ASharp focussharp_focus_optical_table.pysharp_focus_optimizer.pyN/A★ The large-scale setup functions are defined inxl_optical_table.pyandxl_optimizer.py.👀 Overview:In this section we list the available functions in different files and a brief description:Inwave_optics.py: module for scalar optical fields.ClassFunctionsDescriptionScalarLightClass for scalar optical fields defined in the XY plane: complex amplitude $U(r) = A(r)*e^{-ikz}$..drawPlots intensity and phase..apply_circular_maskApply a circular mask of variable radius..apply_triangular_maskApply a triangular mask of variable size..apply_rectangular_maskApply a rectangular mask of variable size..apply_annular_apertureApply annular aperture of variable size..RS_propagationRayleigh-Sommerfelddiffraction integral in z-direction (z>0 and z<0)..get_RS_minimum_zGiven a quality factor, determines the minimum (trustworthy) distance forRS_propagation..CZTChirped z-transform- efficient diffraction using the Bluestein method.LightSourceClass for scalar optical fields defined in the XY plane - defines light source beams..gaussian_beamGaussian beam..plane_wavePlane wave.Invectorized_optics.py: module for vectorized optical fields.ClassFunctionsDescriptionVectorizedLightClass for vectorized optical fields defined in the XY plane: $\vec{E} = (E_x, E_y, E_z)$.drawPlots intensity, phase and amplitude..draw_intensity_profilePlots intensity profile..VRS_propagationVectorial Rayleigh-Sommerfelddiffraction integral in z-direction (z>0 and z<0)..get_VRS_minimum_zGiven a quality factor, determines the minimum (trustworthy) distance forVRS_propagation..VCZTVectorized Chirped z-transform- efficient diffraction using the Bluestein method.PolarizedLightSourceClass for polarized optical fields defined in the XY plane - defines light source beams..gaussian_beamGaussian beam..plane_wavePlane wave.Inoptical_elements.py: shelf with all the optical elements available.FunctionDescriptionScalar light devices-phase_scalar_SLMPhase mask for the spatial light modulator available for scalar fields.SLMSpatial light modulator: applies a phase mask to incident scalar field.Jones matrices-jones_LPJones matrix of alinear polarizerjones_general_retarderJones matrix of ageneral retarder.jones_sSLMJones matrix of thesuperSLM.jones_LCDJones matrix of liquid crystal display (LCD).Polarization-based devices-sSLMsuper-Spatial Light Modulator: adds phase mask (pixel-wise) to $E_x$ and $E_y$ independently.LCDLiquid crystal device: builds any linear wave-plate.linear_polarizerLinear polarizer.BSSingle-side coated dielectric beam splitter.high_NA_objective_lensHigh NA objective lens (only forVectorizedLight).VCZT_objective_lensPropagation through high NA objective lens (only forVectorizedLight).General elements-lensTransparent lens of variable size and focal length.circular_maskCircular mask of variable size.triangular_maskTriangular mask of variable size and orientation.rectangular_maskRectangular mask of variable size and orientation.annular_apertureAnnular aperture of variable size.forked_gratingForked grating of variable size, orientation, and topological charge.Pre-built optical setups-building_blockBasic building unit. Consists of asSLM, andLCDlinked viaVRS_propagation.xl_setupOptical table with the large set-up (Fig.6aofour paper).vSTEDOptical table with the vectorial-based STED setup (Fig.3aofour paper).sharp_focusOptical table with the sharp focus of a radially polarized light beam setup (Fig.4aofour paper).general_setupOptical table with the general setup for large-scale discovery (Fig.5aofour paper).Intoolbox.py: file with useful functions.FunctionDescriptionBasic operations-spaceBuilds the space where light is placed.wrap_phaseWraps any phase mask into $[-\pi, \pi]$ range.is_conserving_energyComputes the total intensity from the light source and compares is with the propagated light -Ref.delta_kroneckerKronecker delta.build_LCD_cellBuilds the cell forLCD.draw_sSLMPlots the two phase masks ofsSLM.moving_avgCompute the moving average of a dataset.rotate_maskRotates the (X, Y) frame w.r.t. given point.profileDetermines the profile of a given input without using interpolation.spot_sizeComputes the spot size as  $\pi (FWHM_x \cdot FWHM_y) /\lambda^2$.compute_fwhmComputes FWHM in 2D.Inloss_functions.py: file with loss functions.FunctionDescriptionvMSE_IntensityParallel computation of Mean Squared Error (Intensity) for a given electric field component $E_x$, $E_y$ or $E_z$.MSE_IntensityMean Squared Error (Intensity) for a given electric field component $E_x$, $E_y$ or $E_z$.vMSE_PhaseParallel computation of Mean Squared Error (Phase) for a given electric field component $E_x$, $E_y$ or $E_z$.MSE_PhaseMean Squared Error (Phase) for a given electric field component $E_x$, $E_y$ or $E_z$.vMSE_AmplitudeParallel computation of Mean Squared Error (Amplitude) for a given electric field component $E_x$, $E_y$ or $E_z$.MSE_AmplitudeMean Squared Error (Amplitude) for a given electric field component $E_x$, $E_y$ or $E_z$.mean_batch_MSE_IntensityBatch-basedMSE_Intensity.small_areaFraction of intensity comprised inside the area of a mask.small_area_STEDFraction of intensity comprised inside the area of a mask - STED version.⚠️ Considerations when using XLuminA:By default, JAX usesfloat32precision. If necessary, enablejax.config.update("jax_enable_x64", True)at the beginning of the file.Basic units are microns (um) and radians. Other units (centimeters, millimeters, nanometers, and degrees) are available at__init.py__.IMPORTANT- RAYLEIGH-SOMMERFELD PROPAGATION:FFT-based diffraction calculation algorithmscan be innacurate depending on the computational window size (sampling).Before propagating light, one should check which is the minimum distance available for the simulation to be accurate.You can use the following functions:get_RS_minimum_z, forScalarLightclass, andget_VRS_minimum_z, forVectorizedLightclass.💻 Development:Some functionalities of XLuminA’s optics simulator (e.g., optical propagation algorithms, planar lens or amplitude masks) are inspired in an open-source NumPy-based Python module for diffraction and interferometry simulation,Diffractio, although we have rewritten and modified these approaches to combine them with JAX just-in-time (jit) functionality. On top of that, we developed completely new functions (e.g., beam splitters, LCDs or propagation through high NA objective lens with CZT methods, to name a few) which significantly expand the software capabilities.Clone repository:git clone https://github.com/artificial-scientist-lab/XLuminA.git📝 How to cite XLuminA:If you use this software, please cite as:@misc{rodríguez2023xlumina,
  title={XLuminA: An Auto-differentiating Discovery Framework for Super-Resolution Microscopy}, 
  author={Carla Rodríguez and Sören Arlt and Leonhard Möckl and Mario Krenn},      
  year={2023},      
  eprint={2310.08408},      
  archivePrefix={arXiv},      
  primaryClass={physics.optics}      
}