large-scale discovery optical table for sted microscopy benchmark

Author: CarlaRodriguez96

experiments/hybrid_sted_optical_table.py

@@ -0,0 +1,58 @@
+# Setting the path for XLuminA modules:
+import os
+import sys
+current_path = os.path.abspath(os.path.join('..'))
+dir_path = os.path.dirname(current_path)
+module_path = os.path.join(dir_path)
+if module_path not in sys.path:
+    sys.path.append(module_path)
+
+from xlumina.__init__ import um, nm, cm, mm
+from xlumina.vectorized_optics import *
+from xlumina.optical_elements import hybrid_setup_fluorophores
+from xlumina.loss_functions import vectorized_loss_hybrid
+from xlumina.toolbox import space, softmin
+import jax.numpy as jnp
+
+"""
+Large-scale setup for STED microscopy baseline rediscovery:
+
+3x3 initial setup - light gets detected across 6 detectors. 
+"""
+
+# 1. System specs:
+sensor_lateral_size = 824 # Resolution
+wavelength1 = 650*nm
+wavelength2 = 532*nm
+x_total = 2500*um
+x, y = space(x_total, sensor_lateral_size)
+shape = jnp.shape(x)[0]
+
+# 2. Define the optical functions: two orthogonally polarized beams:
+w0 = (1200*um, 1200*um)  
+ls1 = PolarizedLightSource(x, y, wavelength1)
+ls1.gaussian_beam(w0=w0, jones_vector=(1, 1))
+ls2 = PolarizedLightSource(x, y, wavelength2)
+ls2.gaussian_beam(w0=w0, jones_vector=(1, 1))
+
+# 3. Define the output (High Resolution) detection:
+x_out, y_out = jnp.array(space(10*um, 400))
+
+# 4. High NA objective lens specs:
+NA = 0.9 
+radius_lens = 3.6*mm/2 
+f_lens = radius_lens / NA
+
+# 5. Static parameters - don't change during optimization:
+fixed_params = [radius_lens, f_lens, x_out, y_out]
+
+# 6. Define the loss function:
+@jit
+def loss_hybrid_sted(parameters):
+    # Output from hybrid_setup is jnp.array(6, N, N): for 6 detectors
+    i_effective = hybrid_setup_fluorophores(ls1, ls2, ls1, ls2, ls1, ls2, parameters, fixed_params, distance_offset = 10) 
+    
+    # Get the minimum value within loss value array of shape (6, 1, 1) 
+    loss_val = softmin(vectorized_loss_hybrid(i_effective)) 
+    
+    return loss_val