new 6x6 topological discovery optical table

Author: CarlaRodriguez96

experiments/six_times_six_ansatz_with_fixed_PM.py

@@ -0,0 +1,62 @@
+# 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 __init__ import um, nm, cm, mm
+from xlumina.vectorized_optics import *
+from xlumina.optical_elements import six_times_six_ansatz
+from xlumina.loss_functions import vectorized_loss_hybrid
+from xlumina.toolbox import space, softmin
+import jax.numpy as jnp
+
+"""
+Pure topological discovery within 6x6 ansatz for Dorn, Quabis and Leuchs (2004)
+"""
+
+# 1. System specs:
+sensor_lateral_size = 824  # Resolution
+wavelength_1 = 635.0*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, wavelength_1)
+ls1.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)) 
+X, Y = jnp.meshgrid(x,y)
+
+# 4. High NA objective lens specs:
+NA = 0.9 
+radius_lens = 3.6*mm/2 
+f_lens = radius_lens / NA
+
+# 4.1 Fixed phase masks:
+# Polarization converter in Dorn, Quabis, Leuchs (2004):
+pi_half = (jnp.pi - jnp.pi/2) * jnp.ones(shape=(sensor_lateral_size // 2, sensor_lateral_size // 2))
+minus_pi_half = - jnp.pi/2 * jnp.ones(shape=(sensor_lateral_size // 2, sensor_lateral_size // 2))
+PM1_1 = jnp.concatenate((jnp.concatenate((minus_pi_half, pi_half), axis=1), jnp.concatenate((minus_pi_half, pi_half), axis=1)), axis=0)
+PM1_2 = jnp.concatenate((jnp.concatenate((minus_pi_half, minus_pi_half), axis=1), jnp.concatenate((pi_half, pi_half), axis=1)), axis=0)
+
+# Linear grating
+PM2_1 = jnp.sin(2*jnp.pi * Y/1000) * jnp.pi
+PM2_2 = jnp.sin(2*jnp.pi * X/1000) * jnp.pi
+
+# 5. Static parameters - don't change during optimization:
+fixed_params = [radius_lens, f_lens, x_out, y_out, PM1_1, PM1_2, PM2_1, PM2_2]
+
+# 6. Define the loss function:
+def loss_hybrid_fixed_PM(parameters):
+    # Output from hybrid_setup is jnp.array(12, N, N): for 12 detectors
+    i_effective = six_times_six_ansatz(ls1, ls1, ls1, ls1, ls1, ls1, ls1, ls1, ls1, ls1, ls1, ls1, parameters, fixed_params, distance_offset = 9.5)
+    # Get the minimum value within loss value array of shape (12, 1, 1) 
+    loss_val = softmin(vectorized_loss_hybrid(i_effective)) 
+    return loss_val