several modification introduced: fixed dimensions to micrometers, introduced ego loop

2026-10-14 23:26:38 +00:00 · 2026-09-21 11:42:22 +02:00
parent 5f88780426
commit 31859271da
4 changed files with 776 additions and 209 deletions
@@ -56,10 +56,16 @@ def make_1Dplot_of(x_axis,y_axis,plot_zoom_x,save_name="plot_1D",) -> None:
        raise ValueError("Lists x_axis and y_axis must have the same length.")
    plt.figure()
    for i in range(len(x_axis)):
-        plt.plot(x_axis[i], y_axis[i])
+        if x_axis[i] is None:
+            plt.plot(y_axis[i])
+        else:
+            plt.plot(x_axis[i], y_axis[i])
    plt.xlabel("x [um]")
    plt.tight_layout()
-    plt.xlim(-plot_zoom_x, plot_zoom_x)
+    if plot_zoom_x is None:
+        pass
+    else:
+        plt.xlim(-plot_zoom_x, plot_zoom_x)
    plt.savefig(save_name + ".pdf")
    plt.close()

@@ -68,7 +74,7 @@ class SpaceGrid:
    Class representing the spatial grid for the simulation.
    Spatial dimensions are intended to be measured in micrometers [um].
    '''
-    def __init__(self, side_pixels: int, unit_cell_pitch: float, prop_distances: np.ndarray, materials_refr_index: np.ndarray, xy_symmetry=False):
+    def __init__(self, side_pixels: int, unit_cell_pitch: float, materials_refr_index: np.ndarray, prop_distances: np.ndarray = None, xy_symmetry=False):
        self.Nx = side_pixels#for the moment we consider only square grids
        self.Ny = side_pixels
        self.unit_cell_pitch = unit_cell_pitch * 1e-6
@@ -89,18 +95,14 @@ class SpaceGrid:
        self.X, self.Y = np.meshgrid(self.xs, self.ys)
        self.rho = np.sqrt(self.X**2 + self.Y**2)
        self.sampling_period = self.unit_cell_pitch
-        self.d = prop_distances * 1e-6
        
-        self.nm = 1e-9
-        self.m_to_nm = 1e9
-        self.m = 1
-        self.um = 1e-6
-        self.m_to_um = 1e6
-        try:
-            if materials_refr_index.shape != self.d.shape:
-                raise ValueError("Materials array must have the same shape as d array.")
-        except AttributeError:
-            raise ValueError("Materials must be an Numpy array.")
+        if prop_distances is None:
+            self.propagation_distances_set = False
+            self.d = None
+        else:
+            self.propagation_distances_set = True
+            self.d = prop_distances * 1e-6
+        
        self.materials_refr_index = materials_refr_index

    def get_grid_shape(self):
@@ -161,18 +163,20 @@ class SpaceGrid:
        return self.rho
    
 class Field:
-    def __init__(self, input_field: np.ndarray, wavelength: float, space_grid: SpaceGrid, center: np.ndarray = np.array([0,0])):
+    def __init__(self, input_field: np.ndarray, wavelength: float, center: np.ndarray = np.array([0,0])):
        self.center = center * 1e-6
-        self.space_grid = space_grid
-        self.input_field = self._shift_field(input_field)
+        self.input_field = input_field
        self.wavelength = wavelength * 1e-6

-        if self.input_field.shape != self.space_grid.get_grid_shape():
-            raise ValueError("Input field must have the same shape as the space grid it is applied onto.")
-
-    def get_field(self, flatten=False):
+    def get_field(self, space_grid: SpaceGrid = None, flatten=False):
+        if space_grid is None:
+            raise ValueError("A SpaceGrid must be provided when invoking get_field()")
+        self.space_grid = space_grid
        if flatten:
-            return self.input_field.flatten()
+            return self._shift_field(self.input_field).flatten()
+        return self._shift_field(self.input_field)
+
+    def get_field_arr(self):
        return self.input_field

    def get_wavelength(self):
@@ -188,16 +192,17 @@ class Field:
        return shift_field_y

    def plot2D(self, plot_zoom_x:float, plot_zoom_y:float, choose_quantity="amplitude", save_name="results/plot_out"):
-        make_2Dplot_of(self.input_field, 
-                       self.space_grid.X, 
-                       self.space_grid.Y,
-                       plot_zoom_x, 
-                       plot_zoom_y, 
-                       choose_quantity, 
-                       save_name)
+        return 0 #broken atm
+        #make_2Dplot_of(self.input_field, 
+        #               self.space_grid.X, 
+        #               self.space_grid.Y,
+        #               plot_zoom_x, 
+        #               plot_zoom_y, 
+        #               choose_quantity, 
+        #               save_name)

 class Source(Field):
-    def __init__(self, input_field: np.ndarray | str, wavelength: float, space_grid: SpaceGrid, field_norm = False, center: np.ndarray = np.array([0,0])):
+    def __init__(self, input_field: np.ndarray | str, wavelength: float, field_norm = False, center: np.ndarray = np.array([0,0])):
        if isinstance(input_field, str):
            input_field_to_constr = np.load(input_field)
            input_field_shape = input_field_to_constr.shape
@@ -208,7 +213,7 @@ class Source(Field):
        if field_norm:
            input_field_to_constr = input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2))

-        super().__init__(input_field_to_constr, wavelength, space_grid, center=center)
+        super().__init__(input_field_to_constr, wavelength, center=center)
        self.type = "source"

    def get_type(self):
@@ -241,7 +246,6 @@ class Target(Field):
    def __init__(self, input_field: np.ndarray | str | Field ,
                 associated_source: Source, 
                 wavelength: float, 
-                 space_grid: SpaceGrid, 
                 field_norm = False, 
                 center = np.array([0,0])):
        if isinstance(input_field, str):
@@ -254,7 +258,7 @@ class Target(Field):
        if field_norm:
            input_field_to_constr = input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2))

-        super().__init__(input_field_to_constr, wavelength, space_grid, center=center)
+        super().__init__(input_field_to_constr, wavelength, center=center)
        self.type = "target"

        if associated_source.get_wavelength() != self.get_wavelength():
@@ -300,48 +304,24 @@ class PlotArray():
            raise ValueError("Input field must be 1D for plot1D.")

 class Simulation:
-    def  __init__(self, sources:list[Source], targets:list[Target], central_wavelength: float, space_grid: SpaceGrid, n_cpu: int = 1):
+    def  __init__(self, sources:list[Source], 
+                  targets:list[Target], 
+                  central_wavelength: float, 
+                  space_grid: SpaceGrid = None, 
+                  n_cpu: int = 1):
+
        self.sources = sources
        self.targets = targets
-        self.space_grid = space_grid
-        self.central_wavelength = central_wavelength * 1e6 #central_wavelength parameter has little sense to exist. Revise
-
+        self.eps_0 = 8.854e-12
+        self.mu_0 = 4*np.pi*1e-7
+        self.central_wavelength = central_wavelength * 1e-6 #central_wavelength parameter has little sense to exist. Revise
        self.operating_wavelength = self.central_wavelength #initially, the operating_wavelength is the central wavelength
-        self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()]
-        
-        self.ks = 2 * np.pi / self.space_grid.get_sampling_period()
-        self.kappas = np.arange(-self.space_grid.get_Nx()//2+1, self.space_grid.get_Nx()//2+1) * self.ks / self.space_grid.get_Nx()
-        self.KX, self.KY = np.meshgrid(self.kappas, self.kappas)
-        self.K_parallel = np.sqrt(self.KX**2 + self.KY**2)
-        self.nu_parallel = self.K_parallel / (2 * np.pi)
-
-        self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()]
-        self.k_vectors = [2 * np.pi / lda for lda in self.lda_vector]
-        self.phase_factors = [self.k_vectors[i] * self.space_grid.get_distances()[i] * np.sqrt(1 - (self.lda_vector[i] * self.nu_parallel) ** 2 + 0*1j) for i in range(len(self.k_vectors))]
-
-        self.P = [np.exp(1j * self.phase_factors[i]) for i in range(len(self.phase_factors))]
-        self.P_nat = [np.fft.ifftshift(self.P[i]) for i in range(len(self.P))]
-
-        #phase masks initialization
-        self.w_norm = np.zeros((2*self.space_grid.get_S()))
-        self.w_norm_symm = np.zeros(self.space_grid.get_dof())
-        self.phase_mask = self._phase_given_w(self.w_norm)
-        #initialize intermediate fields
-        self.intermediate_fields = [None, None, None]
-
-        #pyfftw init
        self.n_cpu = n_cpu
+
+        if space_grid is not None:
+            self.set_space_grid(space_grid)
+        
        self._configure_fftw(self.n_cpu)
-        self.fft_input_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
-        self.fft_output_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
-        self.ifft_input_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
-        self.ifft_output_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
-        self.fft_input_array.fill(0)
-        self.fft_output_array.fill(0)
-        self.ifft_input_array.fill(0)
-        self.ifft_output_array.fill(0)
-        self.fft_operator = pyfftw.FFTW(self.fft_input_array, self.fft_output_array, axes=(0,1), direction='FFTW_FORWARD', threads=self.n_cpu, flags=('FFTW_ESTIMATE',))
-        self.ifft_operator = pyfftw.FFTW(self.ifft_input_array, self.ifft_output_array, axes=(0,1), direction='FFTW_BACKWARD', threads=self.n_cpu, flags=('FFTW_ESTIMATE',))
        
        self.is_optimization_set = False
        self.is_optimized_flag = False
@@ -391,30 +371,10 @@ class Simulation:
            self.use_filtering = True
            self.filter_radius = None

-        if self.use_filtering:
+        if conic_filter_radius is not None:
            self.filter_radius = conic_filter_radius * 1e-6
-            filtered_phase_mask_1 = self._phase_filter(self.phase_mask[:self.space_grid.get_S()])
-            filtered_phase_mask_2 = self._phase_filter(self.phase_mask[self.space_grid.get_S():])
-            self.filtered_phase_mask = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2))
        else:
-            self.filtered_phase_mask = self.phase_mask
-
-        mask_upper_bounds = self.norm_phase_max * np.where( self.space_grid.get_X().flatten()**2 + self.space_grid.get_Y().flatten()**2 < (self.ms_max_avail_diameter/2)**2, 1, 0)
-        mask_upper_bounds_full = np.concatenate((mask_upper_bounds, mask_upper_bounds),axis=0)
-        mask_lower_bounds_full = np.zeros_like(mask_upper_bounds_full)
-        
-        if self.space_grid.xy_symmetry:
-            mask_upper_bounds_symm = mask_upper_bounds.reshape((self.space_grid.get_Nx(),self.space_grid.get_Ny()))[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1].flatten()
-            mask_upper_bounds_symm_full = np.concatenate((mask_upper_bounds_symm, mask_upper_bounds_symm),axis=0)
-            mask_lower_bounds_symm_full = np.zeros_like(mask_upper_bounds_symm_full)
-            self.chosen_mask_upper_bounds = mask_upper_bounds_symm_full
-            self.chosen_mask_lower_bounds = mask_lower_bounds_symm_full
-        else:
-            self.chosen_mask_upper_bounds = mask_upper_bounds_full
-            self.chosen_mask_lower_bounds = mask_lower_bounds_full
-
-        self._solver_init()
-
+            self.filter_radius = None

    def get_sources(self):
        return self.sources
@@ -452,9 +412,13 @@ class Simulation:
        '''
        assert new_wavelength > 0, "Wavelength must be positive."
        self.operating_wavelength = new_wavelength
+
        self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()]
+        self.eps_r = (self.space_grid.get_materials_refr_index())**2
        self.k_vectors = [2 * np.pi / lda for lda in self.lda_vector]
+        
        self.phase_factors = [self.k_vectors[i] * self.space_grid.get_distances()[i] * np.sqrt(1 - (self.lda_vector[i] * self.nu_parallel) ** 2 + 0*1j) for i in range(len(self.k_vectors))]
+        
        self.P = [np.exp(1j * self.phase_factors[i]) for i in range(len(self.phase_factors))]
        self.P_nat = [np.fft.ifftshift(self.P[i]) for i in range(len(self.P))]
        self.P_dagger_nat = [self.P_nat[i].T.conj() for i in range(len(self.P_nat))]
@@ -497,10 +461,10 @@ class Simulation:
        is locally initialized around each source and target.
        '''
    
-        source_Efields = [source.get_field(flatten=True) for source in self.sources]
+        source_Efields = [source.get_field(space_grid=self.space_grid, flatten=True) for source in self.sources]
        source_centers = [source.get_center() for source in self.sources]

-        target_Efields = [target.get_field(flatten=True) for target in self.targets]
+        target_Efields = [target.get_field(space_grid=self.space_grid, flatten=True) for target in self.targets]
        target_centers = [target.get_center() for target in self.targets]

        self.w_norm = np.zeros((2*self.space_grid.get_S()))
@@ -546,6 +510,8 @@ class Simulation:
        Returns:
            ndarray: The filtered phase mask.
        '''
+        if self.filter_radius is None:
+            raise ValueError("Filter radius is not set.")

        phase_mask_cos = npa.cos(phase_mask)
        phase_mask_sin = npa.sin(phase_mask)
@@ -565,35 +531,42 @@ class Simulation:

    def _forward_propagate(self, source_object: Source):
        #source field
-        source_to_propagate = source_object.get_field()
+        source_to_propagate = source_object.get_field(space_grid=self.space_grid)

        self.fft_input_array[:,:] = np.fft.ifftshift(source_to_propagate.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())))
        self.fft_operator()
+
        self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[0]
+        
        self.ifft_operator()
+
        #field before MS1
-        self.intermediate_fields[0] = np.fft.fftshift(self.ifft_operator.output_array)
-        input_field = self.intermediate_fields[0].flatten()
+        self.intermediate_electric_fields[0] = np.fft.fftshift(self.ifft_operator.output_array)
+
+        input_field = self.intermediate_electric_fields[0].flatten()

        self.fft_input_array[:,:] = np.fft.ifftshift((input_field * np.exp(1j * self.filtered_phase_mask[:self.space_grid.get_S()])).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())))
        self.fft_operator()
        self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[1] #element wise product
        self.ifft_operator()
+
        #field before MS2
-        self.intermediate_fields[1] = np.fft.fftshift(self.ifft_operator.output_array.copy())
-        #print("Field intensity after first phase mask and propagation: %.4e" % np.sum(np.abs(self.intermediate_fields[0])**2))
+        self.intermediate_electric_fields[1] = np.fft.fftshift(self.ifft_operator.output_array.copy())
+
+        #print("Field intensity after first phase mask and propagation: %.4e" % np.sum(np.abs(self.intermediate_electric_fields[0])**2))

        self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * np.exp(1j * self.filtered_phase_mask[self.space_grid.get_S():].reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))))
        self.fft_operator()
        self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[2] #element wise product
        self.ifft_operator()
-        #output field
-        self.intermediate_fields[2] = np.fft.fftshift(self.ifft_operator.output_array.copy())

-        self.output_field = self.intermediate_fields[-1].flatten()
+        #output field
+        self.intermediate_electric_fields[2] = np.fft.fftshift(self.ifft_operator.output_array.copy())
+
+        self.output_field = self.intermediate_electric_fields[-1].flatten()

        #make a plot of intermediate fields
-        #for i, field in enumerate(self.intermediate_fields):
+        #for i, field in enumerate(self.intermediate_electric_fields):
        #    make_2Dplot_of(field, self.space_grid.get_X(), self.space_grid.get_Y(), plot_zoom_x=60, plot_zoom_y=60, choose_quantity="amplitude", save_name=f"results/intermediate_field_{i+1}_plot")


@@ -620,19 +593,19 @@ class Simulation:
        Phi_2_dagger = np.exp(-1j * self.filtered_phase_mask[self.space_grid.get_S():]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))

        #compute the adjoint source
-        self.fft_input_array[:,:] = np.fft.ifftshift(np.sum(self.intermediate_fields[2] * np.conj(target.get_field())) * target.get_field())
+        self.fft_input_array[:,:] = np.fft.ifftshift(np.sum(self.intermediate_electric_fields[2] * np.conj(target.get_field(space_grid=self.space_grid))) * target.get_field(space_grid=self.space_grid))
        self.fft_operator()
        self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[2]
        self.ifft_operator()
        self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * Phi_2_dagger)

-        grad_C_phi_2 = 2 * np.real(-1j * self.intermediate_fields[1].conj() * np.fft.fftshift(self.fft_input_array)).flatten() 
+        grad_C_phi_2 = 2 * np.real(-1j * self.intermediate_electric_fields[1].conj() * np.fft.fftshift(self.fft_input_array)).flatten() 

        self.fft_operator()
        self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[1] #element wise
        self.ifft_operator()

-        grad_C_phi_1 = 2 * np.real(-1j * self.intermediate_fields[0].conj() * np.fft.fftshift(self.ifft_operator.output_array) * Phi_1_dagger).flatten()
+        grad_C_phi_1 = 2 * np.real(-1j * self.intermediate_electric_fields[0].conj() * np.fft.fftshift(self.ifft_operator.output_array) * Phi_1_dagger).flatten()
        
        return grad_C_phi_1, grad_C_phi_2

@@ -644,11 +617,10 @@ class Simulation:
        wavelength = source.get_wavelength()

        self._update_wavelength(wavelength)
-        #we shall consider using a conic filter: before doing the forward propagation,
-        #we filter the masks
+
        self._forward_propagate(source)

-        C_s = np.real(np.abs(np.sum(self.output_field * np.conj(target.get_field(flatten=True)))) ** 2)
+        C_s = np.real(np.abs(np.sum(self.output_field * np.conj(target.get_field(space_grid=self.space_grid, flatten=True)))) ** 2)
        grad_C_phi_1, grad_C_phi_2 = self._adjoint_propagate(target)

        return C_s, grad_C_phi_1, grad_C_phi_2
@@ -802,6 +774,7 @@ class Simulation:

        #after the phase mask is set, we test the same mask across all targets
        C_all_wavelengths = []
+        tested_wavelengths = []
        grad_partial = np.zeros_like(grad)

        for target in self.targets:
@@ -809,7 +782,7 @@ class Simulation:
            C, full_grad_1, full_grad_2 = self._value_and_grad(target) 

            C_all_wavelengths.append(C)
-
+            tested_wavelengths.append(target.get_wavelength())
            if self.use_filtering:
                jac_grad_1 = tensor_jacobian_product(self._phase_filter)(phase_mask_1, full_grad_1)
                jac_grad_2 = tensor_jacobian_product(self._phase_filter)(phase_mask_2, full_grad_2)
@@ -831,7 +804,19 @@ class Simulation:

        #the new cost function considers the average overlap across all wavelengths.
        self.opt_history.append(np.mean(C_all_wavelengths))
+
        self._plot_history()
+
+        if self.iter_num[0] % 5 == 0 and self.results_while_solving:
+            fig, ax = plt.subplots()
+            ax.plot(np.array(tested_wavelengths)*1e6, C_all_wavelengths, 'o-')
+            ax.set_xlabel('Wavelength [um]')
+            ax.set_ylabel('Overlap')
+            ax.set_title('Overlap Across Wavelengths')
+            fig.tight_layout()
+            fig.savefig(f"{self.results_folder}/overlap_across_wavelengths.pdf")
+            plt.close(fig)
+
        return np.mean(C_all_wavelengths) 

    def _plot_history(self):
@@ -854,7 +839,29 @@ class Simulation:
        This method sets up the optimizer based on the chosen strategy.

        '''
+        if self.use_filtering:
+            filtered_phase_mask_1 = self._phase_filter(self.phase_mask[:self.space_grid.get_S()])
+            filtered_phase_mask_2 = self._phase_filter(self.phase_mask[self.space_grid.get_S():])
+            self.filtered_phase_mask = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2))
+        else:
+            self.filtered_phase_mask = self.phase_mask
+
+        mask_upper_bounds = self.norm_phase_max * np.where( self.space_grid.get_X().flatten()**2 + self.space_grid.get_Y().flatten()**2 < (self.ms_max_avail_diameter/2)**2, 1, 0)
+        mask_upper_bounds_full = np.concatenate((mask_upper_bounds, mask_upper_bounds),axis=0)
+        mask_lower_bounds_full = np.zeros_like(mask_upper_bounds_full)
+        
+        if self.space_grid.xy_symmetry:
+            mask_upper_bounds_symm = mask_upper_bounds.reshape((self.space_grid.get_Nx(),self.space_grid.get_Ny()))[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1].flatten()
+            mask_upper_bounds_symm_full = np.concatenate((mask_upper_bounds_symm, mask_upper_bounds_symm),axis=0)
+            mask_lower_bounds_symm_full = np.zeros_like(mask_upper_bounds_symm_full)
+            self.chosen_mask_upper_bounds = mask_upper_bounds_symm_full
+            self.chosen_mask_lower_bounds = mask_lower_bounds_symm_full
+        else:
+            self.chosen_mask_upper_bounds = mask_upper_bounds_full
+            self.chosen_mask_lower_bounds = mask_lower_bounds_full
+
        if self.strategy == 'mean':
+
            self.solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof())

            self.solver.set_lower_bounds(self.chosen_mask_lower_bounds.flatten())
@@ -864,7 +871,9 @@ class Simulation:
            self.solver.set_maxeval(self.opt_max_eval)

        elif self.strategy == 'epigraph':
+
            self.solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof() + 1)
+
            epigraph_and_weights_lower_bounds = np.insert(self.chosen_mask_lower_bounds.flatten(), 0, -np.inf)
            epigraph_and_weights_upper_bounds = np.insert(self.chosen_mask_upper_bounds.flatten(), 0, np.inf)

@@ -894,11 +903,20 @@ class Simulation:
        Placeholder method for exporting the results of the optimization.
        This method should be implemented to save the relevant data.
        '''
+        #intermediate fields
+        #intermediate_field[0] is on MS1
+        #intermediate_field[1] is on MS2
+        #intermediate_field[2] is the output
+        for i,field in enumerate(self.intermediate_electric_fields):
+            np.save(f"{self.results_folder}/intermediate_electric_fields_{i}.npy", field)
+            make_2Dplot_of(field, self.space_grid.get_X(), self.space_grid.get_Y(), plot_zoom_x=1e6*self.get_ms_max_avail_diameter()/2, plot_zoom_y=1e6*self.get_ms_max_avail_diameter()/2, choose_quantity="amplitude", save_name=f"{self.results_folder}/intermediate_electric_fields_{i}_plot")
+            make_2Dplot_of(field, self.space_grid.get_X(), self.space_grid.get_Y(), plot_zoom_x=1e6*self.get_ms_max_avail_diameter()/2, plot_zoom_y=1e6*self.get_ms_max_avail_diameter()/2, choose_quantity="phase", save_name=f"{self.results_folder}/intermediate_electric_fields_{i}_phase_plot")
+
        #output fields
        output_field = []
        for i,source in enumerate(self.sources):
            self._forward_propagate(source)
-            output_field.append(self.intermediate_fields[-1])
+            output_field.append(self.intermediate_electric_fields[-1])

            np.save(f"{self.results_folder}/output_fields_{i}.npy", output_field[-1])
            make_2Dplot_of(output_field[-1], self.space_grid.get_X(), self.space_grid.get_Y(), plot_zoom_x=1e6*self.get_ms_max_avail_diameter()/2, plot_zoom_y=1e6*self.get_ms_max_avail_diameter()/2, choose_quantity="amplitude", save_name=f"{self.results_folder}/output_fields_{i}_plot")
@@ -938,78 +956,144 @@ class Simulation:
    def is_optimized(self):
        return self.is_optimized_flag

-    def update_space_grid(self, new_space_grid: SpaceGrid):
+    def set_space_grid(self, new_space_grid: SpaceGrid):
+        '''
+        Update the space grid and associate parameters used for propagation.
+        '''
+        for source in self.sources:
+            if source.get_field_arr().shape != new_space_grid.get_grid_shape():
+                raise ValueError("The provided space grid is not compatible with the source fields.")
+        for target in self.targets:
+            if target.get_field_arr().shape != new_space_grid.get_grid_shape():
+                raise ValueError("The provided space grid is not compatible with the target fields.")
+
+        self.space_grid = new_space_grid
+
+        self.space_grid_is_set = False
        self.is_optimized_flag = False
        self.best_opt_val = 0
-        self.space_grid = new_space_grid
-        self.ks = 2 * np.pi / self.space_grid.get_sampling_period()
-        self.kappas = np.arange(-self.space_grid.get_Nx()//2+1, self.space_grid.get_Nx()//2+1) * self.ks / self.space_grid.get_Nx()
-        self.KX, self.KY = np.meshgrid(self.kappas, self.kappas)
-        self.K_parallel = np.sqrt(self.KX**2 + self.KY**2)
-        self.nu_parallel = self.K_parallel / (2 * np.pi)
+
+        if self.space_grid.get_distances() is not None:
+            try:
+                if self.space_grid.get_materials_refr_index().shape != self.space_grid.get_distances().shape:
+                    raise ValueError("Materials array must have the same shape as d array.")
+            except AttributeError:
+                raise ValueError("Materials must be an Numpy array.")
+            self.space_grid_is_set = True
+            self.d = self.space_grid.get_distances()
+            print("Using distances: ", self.d)
+            self.ks = 2 * np.pi / self.space_grid.get_sampling_period()
+            self.kappas = np.arange(-self.space_grid.get_Nx()//2+1, self.space_grid.get_Nx()//2+1) * self.ks / self.space_grid.get_Nx()
+            self.KX, self.KY = np.meshgrid(self.kappas, self.kappas)
+            self.K_parallel = np.sqrt(self.KX**2 + self.KY**2)
+            self.nu_parallel = self.K_parallel / (2 * np.pi)
+            
+            #phase masks initialization
+            self.w_norm = np.zeros((2*self.space_grid.get_S()))
+            self.w_norm_symm = np.zeros(self.space_grid.get_dof())
+            self.phase_mask = self._phase_given_w(self.w_norm)
+            #initialize intermediate fields
+            self.intermediate_electric_fields = [None, None, None]
+            self.intermediate_magnetic_fields = [None, None, None]
+            #pyfftw init
+            self.fft_input_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
+            self.fft_output_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
+            self.ifft_input_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
+            self.ifft_output_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16)
+            self.fft_input_array.fill(0)
+            self.fft_output_array.fill(0)
+            self.ifft_input_array.fill(0)
+            self.ifft_output_array.fill(0)
+            self.fft_operator = pyfftw.FFTW(self.fft_input_array, self.fft_output_array, axes=(0,1), direction='FFTW_FORWARD', threads=self.n_cpu, flags=('FFTW_ESTIMATE',))
+            self.ifft_operator = pyfftw.FFTW(self.ifft_input_array, self.ifft_output_array, axes=(0,1), direction='FFTW_BACKWARD', threads=self.n_cpu, flags=('FFTW_ESTIMATE',))
+        else:
+            self.space_grid_is_set = False
+
+    def update_propagation_distances(self, new_distances: np.ndarray):
+        '''
+        Update the propagation distances of the simulation by creating a new space grid 
+        with the specified distances while maintaining the other properties of the current space grid.
+        '''
+        self.opt_history = [] #updating distances resets the history
+        new_space_grid = SpaceGrid(
+                side_pixels=self.space_grid.get_Nx(), 
+                unit_cell_pitch=self.space_grid.get_unit_cell_pitch() * 1e6, 
+                prop_distances=new_distances, 
+                materials_refr_index=self.space_grid.get_materials_refr_index(),
+                xy_symmetry=True,
+            )
+        
+        self.set_space_grid(new_space_grid)

    def run(self):
-        '''
-        Placeholder method for running the simulation.
-        This method should be implemented with the actual simulation logic.
-        '''
-        if self.is_optimization_set:
-            print("Starting optimization...")
-            self._initialize_phase_masks()
-            #initialize nlopt solver
+        """
+        Run the optimization process for the simulation.
+        This method checks if the space grid and optimization settings are set,
+        initializes the phase masks, and performs the optimization based on the chosen strategy.
+        """
+        if self.space_grid_is_set:
+            if self.is_optimization_set:
+                print("Initializing solver...")
+                self._solver_init()
+                print("Starting optimization...")
+                self._initialize_phase_masks()
+                #initialize nlopt solver

-            if self.space_grid.xy_symmetry:
-                if self.strategy == 'epigraph':
-                    epigraph_w_norm_symm = self.solver.optimize(np.insert(self.w_norm_symm, 0, 0))
-                    self.w_norm_symm[:] = epigraph_w_norm_symm[1:]
-                elif self.strategy == 'mean':
-                    self.w_norm_symm[:] = self.solver.optimize(self.w_norm_symm)
+                if self.space_grid.xy_symmetry:
+                    if self.strategy == 'epigraph':
+                        epigraph_w_norm_symm = self.solver.optimize(np.insert(self.w_norm_symm, 0, 0))
+                        self.w_norm_symm[:] = epigraph_w_norm_symm[1:]
+                    elif self.strategy == 'mean':
+                        self.w_norm_symm[:] = self.solver.optimize(self.w_norm_symm)
+                    else:
+                        raise ValueError(f"Unknown strategy: {self.strategy}")
+
+                    self.is_optimized_flag = True
+                    self.best_opt_val = self.solver.last_optimum_value()
+                    self.w_norm[:self.space_grid.get_S()] = self._expand_symmetrize(self.w_norm_symm[:self.space_grid.get_dof()//2])
+                    self.w_norm[self.space_grid.get_S():] = self._expand_symmetrize(self.w_norm_symm[self.space_grid.get_dof()//2:])
+                    self.phase_mask[:] = self._phase_given_w(self.w_norm)
+
+                    if self.use_filtering:
+                        phase_mask_1 = self.phase_mask[:self.space_grid.get_S()]
+                        phase_mask_2 = self.phase_mask[self.space_grid.get_S():]
+                        filtered_phase_mask_1 = self._phase_filter(phase_mask_1)
+                        filtered_phase_mask_2 = self._phase_filter(phase_mask_2)
+                        self.filtered_phase_mask[:] = np.concatenate([filtered_phase_mask_1, filtered_phase_mask_2])
+                    else:
+                        self.filtered_phase_mask[:] = self.phase_mask
+                    
+                    self.is_optimization_set = True
+                    self.best_opt_val = self.solver.last_optimum_value()
                else:
-                    raise ValueError(f"Unknown strategy: {self.strategy}")
-                self.is_optimized_flag = True
-                self.best_opt_val = self.solver.last_optimum_value()
-                self.w_norm[:self.space_grid.get_S()] = self._expand_symmetrize(self.w_norm_symm[:self.space_grid.get_dof()//2])
-                self.w_norm[self.space_grid.get_S():] = self._expand_symmetrize(self.w_norm_symm[self.space_grid.get_dof()//2:])
-                self.phase_mask[:] = self._phase_given_w(self.w_norm)
+                    if self.strategy == 'epigraph':
+                        self.epigraph_w_norm = np.insert(self.w_norm, 0, 0)
+                        self.epigraph_w_norm[:] = self.solver.optimize(self.epigraph_w_norm)
+                        self.w_norm[:] = self.epigraph_w_norm[1:]
+                    elif self.strategy == 'mean':
+                        self.w_norm[:] = self.solver.optimize(self.w_norm)
+                    else:
+                        raise ValueError(f"Unknown strategy: {self.strategy}")
+                    self.phase_mask[:] = self._phase_given_w(self.w_norm)
                    
-                if self.use_filtering:
-                    phase_mask_1 = self.phase_mask[:self.space_grid.get_S()]
-                    phase_mask_2 = self.phase_mask[self.space_grid.get_S():]
-                    filtered_phase_mask_1 = self._phase_filter(phase_mask_1)
-                    filtered_phase_mask_2 = self._phase_filter(phase_mask_2)
-                    self.filtered_phase_mask[:] = np.concatenate([filtered_phase_mask_1, filtered_phase_mask_2])
-                else:
-                    self.filtered_phase_mask[:] = self.phase_mask
+                    if self.use_filtering:
+                        phase_mask_1 = self.phase_mask[:self.space_grid.get_S()]
+                        phase_mask_2 = self.phase_mask[self.space_grid.get_S():]
+                        filtered_phase_mask_1 = self._phase_filter(phase_mask_1)
+                        filtered_phase_mask_2 = self._phase_filter(phase_mask_2)
+                        self.filtered_phase_mask[:] = np.concatenate([filtered_phase_mask_1, filtered_phase_mask_2])
+                    else:
+                        self.filtered_phase_mask[:] = self.phase_mask
                    
-                self.is_optimization_set = True
-                self.best_opt_val = self.solver.last_optimum_value()
+                    self.is_optimized_flag = True
+                    self.best_opt_val = self.solver.last_optimum_value()
+
+                if self.save_results:
+                    self._export_results()
+
+                self.iter_num[:] = [0]
+                return self.best_opt_val
            else:
-                if self.strategy == 'epigraph':
-                    self.epigraph_w_norm = np.insert(self.w_norm, 0, 0)
-                    self.epigraph_w_norm[:] = self.solver.optimize(self.epigraph_w_norm)
-                    self.w_norm[:] = self.epigraph_w_norm[1:]
-                elif self.strategy == 'mean':
-                    self.w_norm[:] = self.solver.optimize(self.w_norm)
-                else:
-                    raise ValueError(f"Unknown strategy: {self.strategy}")
-                self.phase_mask[:] = self._phase_given_w(self.w_norm)
-                
-                if self.use_filtering:
-                    phase_mask_1 = self.phase_mask[:self.space_grid.get_S()]
-                    phase_mask_2 = self.phase_mask[self.space_grid.get_S():]
-                    filtered_phase_mask_1 = self._phase_filter(phase_mask_1)
-                    filtered_phase_mask_2 = self._phase_filter(phase_mask_2)
-                    self.filtered_phase_mask[:] = np.concatenate([filtered_phase_mask_1, filtered_phase_mask_2])
-                else:
-                    self.filtered_phase_mask[:] = self.phase_mask
-                
-                self.is_optimized_flag = True
-                self.best_opt_val = self.solver.last_optimum_value()
-
-            if self.save_results:
-                self._export_results()
-
-            return self.best_opt_val
+                raise RuntimeError("Optimization settings have not been specified.")
        else:
-            raise RuntimeError("Optimization settings have not been specified.")
-        
+            raise RuntimeError("Space grid has not been specified.")
@@ -5,20 +5,23 @@ if not os.path.exists("results"):
    os.makedirs("results")
 #Space grid definition
 N = 1025
-unit_cell_pitch = 700e-9
+unit_cell_pitch = 0.700

 n0 = 1.45
 n1 = 1.0
 n2 = 1.45
 mat_refr_index = np.array([n0, n1, n2])

-d0 = 476e-6 #propagation distance: source - MS1 [m]
-d1 = 510e-6 #propagation distance: MS1 - MS2 [m]
-d2 = d0 #propagation distance: MS2 - target [m]
+
+#d0 = 400 #propagation distance: source - MS1 [um] NO FILT
+#d1 = 280 #propagation distance: MS1 - MS2 [um] NO FILT
+d0 = 535 #propagation distance: source - MS1 [um] FILTERED
+d1 = 460 #propagation distance: MS1 - MS2 [um] FILTERED
+d2 = d0 #propagation distance: MS2 - target [um]
 d = np.array([d0, d1, d2])

-ms_max_available_diameter = 122e-6 #[um]
-output_fiber_core_diameter = 19e-6 #[um]
+ms_max_available_diameter = 122 #[um]
+output_fiber_core_diameter = 19 #[um]
 opt_max_eval = 140

 space_grid = pid.SpaceGrid(
@@ -30,38 +33,39 @@ space_grid = pid.SpaceGrid(
 )

 #define the source
-beam_waist = 5.2e-6
+beam_waist = 5.2
 gaussian_field = np.exp(-(space_grid.get_rho())**2 / (beam_waist)**2) #Gaussian input field, flattened
 gaussian_field_normalized = (gaussian_field / np.sqrt(np.sum(np.abs(gaussian_field)**2))) #normalize the input field to have power = 1

-wavelengths_to_test_nm = [1550] #[nm]
-central_wavelength_m = 1550 * space_grid.nm
-sources = [pid.Source(input_field=gaussian_field_normalized, 
+wavelengths_to_test_um = [1.530,1.550,1.570] #[um]
+source_files = ["fiber_modes/LP01_field_distribution_lda_1530_N_1025.npy",
+                "fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy",
+                "fiber_modes/LP01_field_distribution_lda_1570_N_1025.npy"]
+sources = [pid.Source(input_field=np.load(file), 
                      wavelength=wvl, 
-                      space_grid=space_grid, 
-                      wvl_unit_measure=space_grid.nm) for wvl in wavelengths_to_test_nm]
+                      field_norm=True) for file, wvl in zip(source_files, wavelengths_to_test_um)]

-sources[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/input_field')
+#sources[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/input_field')
 #define targets
 #import field data from files and create target objects
-target_files = ["fiber_modes/LP21_field_distribution_lda_1550_N_1025.npy",]
+target_files = ["fiber_modes/LP21_field_distribution_lda_1530_N_1025.npy",
+                "fiber_modes/LP21_field_distribution_lda_1550_N_1025.npy",
+                "fiber_modes/LP21_field_distribution_lda_1570_N_1025.npy"]

 target_fields = [np.load(file) for file in target_files]

 targets = [pid.Target(input_field=target_fields[i], 
                      associated_source=sources[i], 
                      wavelength=wvl, 
-                      space_grid=space_grid, 
-                      wvl_unit_measure=space_grid.nm,
-                      field_norm = True) for i, (file, wvl) in enumerate(zip(target_files, wavelengths_to_test_nm))]
+                      field_norm = True) for i, (file, wvl) in enumerate(zip(target_files, wavelengths_to_test_um))]

-targets[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/target_field')
+#targets[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/target_field')

 #setup the simulation instance
 simulation = pid.Simulation(
    sources=sources,
    targets=targets,
-    central_wavelength=central_wavelength_m,
+    central_wavelength=wavelengths_to_test_um[0],
    space_grid=space_grid,
    n_cpu=1
 )
@@ -72,7 +76,7 @@ simulation.set_optimization(
    ms_max_avail_diameter=ms_max_available_diameter,
    phase_mask_init_size=output_fiber_core_diameter/3,
    opt_max_eval=opt_max_eval,
-    conic_filter_radius=1.5e-6,    
+    conic_filter_radius=1.0,    
 )

 simulation.run()
@@ -0,0 +1,479 @@
+import mpid as pid
+import numpy as np
+from mpi4py.futures import MPIPoolExecutor
+
+from smt.applications.ego import EGO, KRG
+from smt.sampling_methods import LHS
+from smt.design_space import (
+    DesignSpace,
+)
+from sklearn.model_selection import KFold, LeaveOneOut
+from scipy.stats import norm
+from scipy.optimize import minimize
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+import matplotlib.animation as animation
+import copy
+import pickle
+import os
+import gc
+if not os.path.exists("results"):
+    os.makedirs("results")
+#Space grid definition
+N = 1025
+unit_cell_pitch = 0.700 #i.e., 700 [nm]
+
+n0 = 1.45
+n1 = 1.0
+n2 = 1.45
+mat_refr_index = np.array([n0, n1, n2])
+
+ms_max_available_diameter = 122 #[um]
+output_fiber_core_diameter = 19 #[um]
+opt_max_eval = 140
+
+def setup_optical_system()->pid.Simulation:
+    '''
+    Builds the mode converter optical system using 
+    the specified propagation distances `d`, runs the optimization
+    and returns the best objective value.
+
+    Args:
+        xdoe_points (np.ndarray): Array containing the DOE points.
+        Each point consists of the propagation distances d0 and d1
+
+    Returns:
+        float: The best objective value obtained from the optimization.
+    '''
+    central_wavelength_um = 1.55
+    wavelength = central_wavelength_um
+
+    wavelengths_to_test_um = [1.530, 1.550, 1.570] #[um]
+
+    source_files = ["fiber_modes/LP01_field_distribution_lda_1530_N_1025.npy",
+                    "fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy",
+                    "fiber_modes/LP01_field_distribution_lda_1570_N_1025.npy"]
+
+    # source = pid.Source(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy', 
+    #                     wavelength=wavelength, 
+    #                     field_norm=True)
+
+    sources = [pid.Source(input_field=np.load(file), 
+                          wavelength=wvl, 
+                          field_norm=True) for file, wvl in zip(source_files, wavelengths_to_test_um)]
+
+    target_files = ["fiber_modes/LP21_field_distribution_lda_1530_N_1025.npy",
+                    "fiber_modes/LP21_field_distribution_lda_1550_N_1025.npy",
+                    "fiber_modes/LP21_field_distribution_lda_1570_N_1025.npy"]
+
+    targets = [pid.Target(input_field=np.load(file), 
+                          associated_source=sources[i], 
+                          wavelength=wvl, 
+                          field_norm=True) for i, (file, wvl) in enumerate(zip(target_files, wavelengths_to_test_um))]
+    
+    # target = pid.Target(input_field='fiber_modes/LP21_field_distribution_lda_1550_N_1025.npy', 
+    #                     associated_source=source, 
+    #                     wavelength=wavelength, 
+    #                     field_norm=True)
+  
+    initial_space_grid = pid.SpaceGrid(
+        side_pixels=N, 
+        unit_cell_pitch=unit_cell_pitch, 
+        materials_refr_index=mat_refr_index,
+        xy_symmetry=True,
+    )
+
+    #setup the simulation instance
+    simulation = pid.Simulation(
+        sources=sources,
+        targets=targets,
+        central_wavelength=central_wavelength_um,
+        space_grid=initial_space_grid,
+        n_cpu=1
+    )
+
+    simulation.set_optimization(
+        opt_param_min=0,
+        opt_param_max=1,
+        ms_max_avail_diameter=ms_max_available_diameter,
+        phase_mask_init_size=output_fiber_core_diameter/3,
+        opt_max_eval=opt_max_eval,
+        conic_filter_radius=1.0,
+        results_while_solving = False,
+        save_results = False,
+    )
+
+    return simulation
+
+def EI(GP, points, f_min):
+    pred = - GP.predict_values(points) #GP predicts overlaps (positive)
+    var = GP.predict_variances(points)
+    args0 = (f_min - pred) / np.sqrt(var)
+    args1 = (f_min - pred) * norm.cdf(args0)
+    args2 = np.sqrt(var) * norm.pdf(args0)
+
+    if var.size == 1 and var == 0.0:  # can be use only if one point is computed
+        return 0.0
+
+    ei = args1 + args2
+    return ei[0,0]
+
+def leave_one_out(initial_dataset):
+    full_x_doe = initial_dataset[:, :2]
+    full_y_doe = initial_dataset[:, 2].reshape(-1, 1)
+    loo = LeaveOneOut()
+
+    #available kernels: ‘matern52’, ‘matern32’ ‘abs_exp’, ‘squar_exp’
+    kernels_to_test = ['abs_exp', 'squar_exp', 'matern52', 'matern32']
+    all_z_values = []
+
+    for kernel in kernels_to_test:
+        errs = []
+        z_values = []
+        true_values = []
+        predicted_values = []
+        predicted_confidence = []
+
+        for i, (train_index, test_index) in enumerate(loo.split(full_x_doe)):
+
+            print(f"LOO Iteration {i}")
+
+            train_values = full_x_doe[train_index]
+            test_values = full_x_doe[test_index]
+
+            train_outputs = full_y_doe[train_index]
+            test_outputs = full_y_doe[test_index]
+
+            sm_loo = KRG(design_space=design_space, 
+                        n_start=40, 
+                        print_global=False,
+                        eval_noise=True,
+                        corr=kernel,
+                        )
+            #no_filt has a good fit using matern32
+            #filt has a good fit using matern32
+            sm_loo.set_training_values(train_values, train_outputs)
+            sm_loo.train()
+
+            #this is the predicted value from the surrogate model
+            test_pred_values = sm_loo.predict_values(test_values) 
+            #for each point in the surrogate model we can get an estimate of the prediction variance, in other
+            #words, a measure of how much the model trusts the prediction at that point
+            test_pred_variances = sm_loo.predict_variances(test_values) 
+            #standardized error: this value should roughly behave as a normal
+            #distribution (Gaussian) with mean 0 and variance 1
+            z = (test_outputs - test_pred_values) / np.sqrt(test_pred_variances)
+            #note that we divide for the standard deviation, thus the deviation at 
+            #the numerator shall be within +/- 3 standard deviations, for most points
+            z_values.append(z)
+            true_values.append(test_outputs)
+            predicted_values.append(test_pred_values)
+            predicted_confidence.append(np.sqrt(test_pred_variances)) #standard deviation
+            sq_err = np.mean((test_pred_values - test_outputs)**2)
+
+            print(f" Squared error for LOO iteration {i}: {sq_err:.6e}")
+            errs.append(sq_err)
+        all_z_values.append(z_values)
+
+        fig, ax = plt.subplots(figsize=(10, 6))
+        for i in range(len(true_values)):
+            #confidence interval is set to 3 standard deviations (99.7%)
+            ax.errorbar(i, predicted_values[i], yerr=3 * predicted_confidence[i], fmt='o', label='Predicted with confidence' if i == 0 else "")
+            ax.scatter(i, true_values[i], marker='x', color='red', label='True value' if i == 0 else "")
+        ax.set_xticks(range(len(true_values)))
+        ax.set_xlabel('LOO Iteration', fontsize=20)
+        ax.set_ylabel('Value', fontsize=20)
+        ax.set_title('LOO Predicted vs True Values', fontsize=20)
+        ax.tick_params(axis='both', labelsize=18)
+        ax.grid(True, alpha=0.3)
+        ax.legend()
+        plt.savefig(f'{results_folder}/loo_validation_kernel_{kernel}.pdf')
+        plt.close(fig)
+
+        #make a line plot to show on a diagonal how predicted and true values compare
+        true_values_array = np.concatenate(true_values).ravel()
+        predicted_values_array = np.concatenate(predicted_values).ravel()
+
+        lower = min(true_values_array.min(), predicted_values_array.min())
+        upper = max(true_values_array.max(), predicted_values_array.max())
+
+        fig, ax = plt.subplots(figsize=(10, 6))
+        ax.scatter(true_values_array, predicted_values_array, color="blue")
+        ax.plot([lower, upper], [lower, upper], "r--", label="y=x")
+        ax.set_xlabel('True Value', fontsize=20)
+        ax.set_ylabel('Predicted Value', fontsize=20)
+        ax.set_title('LOO Predicted vs True Values', fontsize=20)
+        ax.tick_params(axis='both', labelsize=18)
+        ax.grid(True, alpha=0.3)
+        ax.legend()
+        plt.savefig(f'{results_folder}/loo_predicted_vs_true_diagonal_kernel_{kernel}.pdf')
+        plt.close(fig)
+
+        fig, ax = plt.subplots(figsize=(10, 6))
+        z_values_array = np.concatenate(z_values)
+        ax.scatter(range(len(z_values_array)), z_values_array, color='blue', label='Standardized residuals')
+        ax.set_ylabel('Standardized Residual', fontsize=20)
+        ax.set_title('LOO Standardized Residuals', fontsize=20)
+        ax.set_xticks(range(len(z_values_array)))
+        ax.set_ylim(min(z_values_array) - 0.5, max(z_values_array) + 0.5)
+        ax.tick_params(axis='both', labelsize=18)
+        ax.grid(True, alpha=0.3)
+        ax.legend()
+        plt.savefig(f'{results_folder}/loo_standardized_residuals_kernel_{kernel}.pdf')
+        plt.close(fig)
+
+    #in the same plot, represent the normal distribution (0,1) and 
+    #all the fitted standardized residuals from the different KRG kernels
+    fig, ax = plt.subplots(figsize=(10, 6))
+    all_std_values = []
+    x = np.linspace(-4, 4, 100)
+    for kernel_name, z_vals in zip(kernels_to_test, all_z_values):
+        mu, std = norm.fit(np.array(z_vals))
+        all_std_values.append(std)
+        print(f"FIT VALUES for kernel {kernel_name}: mu =", mu, "std =", std)
+        p = norm.pdf(x, mu, std)
+        ax.plot(x, p, label=f'ker={kernel_name}')
+
+
+    ax.plot(x, norm.pdf(x, 0, 1), 'g--', label='N(0,1) Gaussian')
+
+    ax.set_xlabel('std. residual', fontsize=20)
+    ax.set_ylabel('Standardized Residual', fontsize=20)
+    ax.tick_params(axis='both', labelsize=18)
+    ax.grid(True, alpha=0.3)
+    ax.legend()
+    plt.savefig(f'{results_folder}/loo_standardized_residuals_gaussian_fit.pdf')
+    plt.close(fig)
+
+    best_kernel = kernels_to_test[np.argmin(np.array(all_std_values))]
+
+    #train the model using the whole dataset
+    sm_final = KRG(design_space=design_space, 
+                   n_start=25, 
+                   print_global=False, 
+                   eval_noise=True, 
+                   corr=best_kernel)
+    sm_final.set_training_values(full_x_doe, full_y_doe)
+    sm_final.train()
+
+    #save the train data
+    with open(f'{surrogate_folder_name}/train_data.pkl', 'wb') as f:
+        pickle.dump((full_x_doe, full_y_doe), f)
+
+    #save the trained surrogate model
+    with open(f'{surrogate_folder_name}/trained_surrogate_model_kernel_{best_kernel}.pkl', 'wb') as f:
+        pickle.dump(sm_final, f)
+    return sm_final
+
+def run_ego(surrogate_model, n_ego_iter, xlimits, simulation):
+    sm = surrogate_model #the surrogate is trained on overlap values (positive)
+    #load the associated fold training data
+    xdoe_train_loaded, ydoe_train_loaded = pickle.load(open(f'{surrogate_folder_name}/train_data.pkl', 'rb'))
+    x_data = copy.deepcopy(xdoe_train_loaded)
+    y_data = - copy.deepcopy(ydoe_train_loaded)
+    n_start = 25 #number of initial random points to sample for EGO refinement
+    x_et_k_tested_history = []
+    EI_history = []
+    for k in range(n_ego_iter):
+        d0_start = np.atleast_2d(xlimits[0][0] + np.random.rand(n_start)*(xlimits[0][1]-xlimits[0][0]))
+        d1_start = np.atleast_2d(xlimits[1][0] + np.random.rand(n_start)*(xlimits[1][1]-xlimits[1][0]))
+        x_start = np.hstack((d0_start.T, d1_start.T))
+        f_min_k = np.min(y_data) #the most negative data point is the best point
+        opt_all = np.array(#maximizes the expected improvement in a region around x_st, an element of the start point
+            [
+                minimize(lambda x: float(-EI(sm, np.atleast_2d(x), f_min_k)), x_st, method="SLSQP", bounds=xlimits)
+                for x_st in x_start
+            ]
+        )
+        opt_success = opt_all[[opt_i["success"] for opt_i in opt_all]]
+        obj_success = np.array([opt_i["fun"] for opt_i in opt_success])
+        ind_min = np.argmin(obj_success) #in this case, the most negative -EI is the best (maximum EI)
+        #save the minimum EI value for this iteration to the history
+        EI_history.append(- obj_success[ind_min])
+        opt = opt_success[ind_min]
+        #select the non-optimal points for history purposes. Later, we will plot
+        #all the non-optimal points with crosses to show the search process of the n_start points
+        non_optimal_points = [opt_i["x"] for i, opt_i in enumerate(opt_all) if i != ind_min]
+        x_et_k = opt["x"] #this is the coordinate of the most interesting point: the one that maximizes the EI (minimize -EI)
+        d0_et_k = x_et_k[0]
+        d1_et_k = x_et_k[1]
+        d2_et_k = d0_et_k
+        d_et_k = np.round(np.array([d0_et_k, d1_et_k, d2_et_k]),0)
+        print(d_et_k)
+        simulation.update_propagation_distances(d_et_k)
+        simulation.run()
+        y_et_k = - simulation.get_best_opt_val() #evaluate the objective function at the new point
+        x_et_k_tested = np.atleast_2d(x_et_k)
+        x_et_k_tested_history.append(x_et_k_tested)
+
+        d0_EI = np.linspace(xlimits[0][0], xlimits[0][1], 100)
+        d1_EI = np.linspace(xlimits[1][0], xlimits[1][1], 100)
+        d0_EI_grid, d1_EI_grid = np.meshgrid(d0_EI, d1_EI)
+        EI_grid = np.zeros_like(d0_EI_grid)
+        for i in range(d0_EI_grid.shape[0]):
+            for j in range(d0_EI_grid.shape[1]):
+                EI_grid[i,j] = EI(sm, np.atleast_2d([d0_EI_grid[i,j], d1_EI_grid[i,j]]), f_min_k)
+
+        fig, ax = plt.subplots(figsize=(8,6))
+        contour = ax.contourf(d0_EI_grid, d1_EI_grid, EI_grid, levels=50, cmap='viridis')
+        fig.colorbar(contour, ax=ax, label='Expected Improvement')
+        #scatter the full history of tested points up to the current iteration k, the current point is red, the others are blue
+        x_et_k_tested_history_array = np.vstack(x_et_k_tested_history)
+        #scatter last point in red, the others in blue
+        ax.scatter(x_et_k_tested_history_array[:-1,0], x_et_k_tested_history_array[:-1,1], color='blue', label='Previous best')
+        ax.scatter(x_et_k_tested_history_array[-1,0], x_et_k_tested_history_array[-1,1], color='red', label='Best point')
+        #also plot the non optimal points with crosses for this iteration
+        if len(non_optimal_points) > 0:
+            non_optimal_points_array = np.vstack(non_optimal_points)
+            ax.scatter(non_optimal_points_array[:,0], non_optimal_points_array[:,1], color='orange', marker='x', label='Candidate')
+        ax.set_title(f'EGO Iteration {k+1}/{n_ego_iter}')
+        ax.set_xlabel('d0 [um]')
+        ax.set_ylabel('d1 [um]')
+        ax.set_xlim(xlimits[0])
+        ax.set_ylim(xlimits[1])
+        ax.legend(loc='upper left', framealpha=0.5)
+        fig.tight_layout()
+        fig.savefig(f'{surrogate_folder_name}/optimisation_{k}.png')
+        plt.close(fig)
+
+        y_data = np.atleast_2d(np.append(y_data, y_et_k)).reshape(-1,1)
+        x_data = np.atleast_2d(np.append(x_data, x_et_k_tested, axis=0))
+        sm.set_training_values(x_data, - y_data) #train the surrogate with positive values.
+        sm.train() #train the model with the new data point added to the previous ones
+        print(f"{k}-th EGO refinement complete")
+        #for this k iteration out of the n_start, make a plot of the EI surface and append the poisition of the 
+        #selected point. Save the figure to a plot list ims, which will be used to create an animation at the end of the EGO loop
+        #To test the EI surface, use a custom grid of d0,d1 points, and compute the EI for each point in the grid. Then plot the EI surface as a contour plot, and mark the selected point with a red dot.
+        
+    ind_best = np.argmin(y_data)
+    x_opt = x_data[ind_best]
+    y_opt = - y_data[ind_best]
+    np.savetxt(f'{surrogate_folder_name}/ego_opt_results.txt', np.hstack((x_opt,y_opt)))
+    ims = []
+    fig = plt.figure(figsize=[10, 10])
+
+    ax = plt.gca()
+    ax.axes.get_xaxis().set_visible(False)
+    ax.axes.get_yaxis().set_visible(False)
+
+    for k in range(n_ego_iter):
+        image_pt = mpimg.imread(f"{surrogate_folder_name}/optimisation_{k}.png")
+        im = plt.imshow(image_pt)
+        ims.append([im])
+
+    ani = animation.ArtistAnimation(fig, ims, interval=500)
+    #save the animation to a gif file
+    ani.save(f'{surrogate_folder_name}/ego_optimization.gif', writer='pillow')
+
+    #plot of the EI history
+
+    fig, ax = plt.subplots(figsize=(8,6))
+    ax.plot(range(1, n_ego_iter+1), EI_history, marker='o')
+    ax.set_title('Expected Improvement History')
+    ax.set_xlabel('EGO Iteration')
+    ax.set_ylabel('Expected Improvement')
+    ax.grid()
+    fig.tight_layout()
+    fig.savefig(f'{surrogate_folder_name}/EI_history.pdf')
+    plt.close(fig)
+
+    #surface plot of the refined surrogate model
+    d0_grid = np.linspace(xlimits[0][0], xlimits[0][1], 100)
+    d1_grid = np.linspace(xlimits[1][0], xlimits[1][1], 100)
+    d0_grid_mesh, d1_grid_mesh = np.meshgrid(d0_grid, d1_grid)
+    d_grid = np.column_stack((d0_grid_mesh.flatten(), d1_grid_mesh.flatten()))
+    sm_predictions = sm.predict_values(d_grid).reshape(d0_grid_mesh.shape)
+    fig = plt.figure(figsize=(8,6))
+    ax = fig.add_subplot(111, projection='3d')
+    surf = ax.plot_surface(d0_grid_mesh, d1_grid_mesh, sm_predictions, cmap='viridis', alpha=0.7, edgecolor='none')
+    ax.scatter(xdoe_train_loaded[:,0], xdoe_train_loaded[:,1], ydoe_train_loaded.flatten(), color='blue', label='DOE Points')
+    ax.scatter(x_data[len(xdoe_train_loaded):,0], x_data[len(xdoe_train_loaded):,1], - y_data[len(xdoe_train_loaded):].flatten(), color='red', label='EGO Evaluated Points')
+    ax.scatter([x_opt[0]], [x_opt[1]], [y_opt], color='green', s=100, label='Best Point', edgecolor='black')
+    ax.set_xlabel('d0 [um]')
+    ax.set_ylabel('d1 [um]')
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(f'{surrogate_folder_name}/refined_surrogate_model_surface.pdf')
+    plt.close(fig)
+    #contour plot
+    fig, ax = plt.subplots(figsize=(8,6))
+    cp = ax.contourf(d0_grid_mesh, d1_grid_mesh, sm_predictions, levels=50, cmap='viridis')
+    fig.colorbar(cp, ax=ax, label='Surrogate Model Prediction')
+    ax.scatter(xdoe_train_loaded[:,0], xdoe_train_loaded[:,1], color='blue', label='DOE Points')
+    ax.scatter(x_data[len(xdoe_train_loaded):,0], x_data[len(xdoe_train_loaded):,1], color='red', label='EGO Evaluated Points')
+    ax.scatter([x_opt[0]], [x_opt[1]], color='green', s=100, label='Best Point', edgecolor='black')
+    ax.set_xlabel('d0 [um]')
+    ax.set_ylabel('d1 [um]')
+    ax.set_xlim(xlimits[0])
+    ax.set_ylim(xlimits[1])
+    ax.legend(framealpha=0.5)
+    fig.tight_layout()
+    fig.savefig(f'{surrogate_folder_name}/refined_surrogate_model_contour.pdf')
+    plt.close(fig)
+
+    pickle.dump(sm, open(f'{surrogate_folder_name}/refined_surrogate_model.pkl', 'wb'))
+    pickle.dump((x_data, y_data), open(f'{surrogate_folder_name}/refined_model_training_data.pkl', 'wb'))
+    return x_opt, y_opt, ind_best, x_data, y_data
+
+
+##############################
+#EGO setup
+##############################
+if __name__ == "__main__":
+    n_ego_iter = 20
+    xlimits = np.array([[100,700], [100,700]]) # d0, d1 in [um]
+    seed = 777
+    np.random.seed(seed)
+    design_space = DesignSpace(xlimits, seed=seed)
+    criterion = 'EI'
+    ndoe = 25 #all the samples will be used for leave-one-out cross-validation
+
+    #first step: check if a DOE exists, if not not, build one
+    surrogate_folder_name = 'single_converter_symm_surrogate_mf_filt'
+    results_folder = 'single_converter_symm_results_ego_mf_filt'
+    if not os.path.exists(surrogate_folder_name):
+        os.makedirs(surrogate_folder_name)
+    if not os.path.exists(results_folder):
+        os.makedirs(results_folder)
+
+    simulation = setup_optical_system()
+
+    if not os.path.exists(f'{surrogate_folder_name}/ego_doe.npy'):
+        #build DOE
+        print("Building DOE...")
+        sampling = LHS(xlimits=xlimits, seed=seed)
+        xdoe = np.round(sampling(ndoe),0)
+        ydoe = []
+
+        ydoe = np.zeros((len(xdoe), 1), dtype=float)
+        
+        for i, point in enumerate(xdoe):
+            print(f"DOE evaluation {i+1}/{len(xdoe)} for point {point}")
+            d0 = point[0]
+            d1 = point[1]
+            d2 = d0
+            d = np.round(np.array([d0, d1, d2]),0)
+            simulation.update_propagation_distances(d)
+            simulation.run()
+            ydoe[i, 0] = simulation.get_best_opt_val()
+
+        with open(f'{surrogate_folder_name}/ego_evaluations.txt', 'a') as f:
+            for xd, yd in zip(xdoe, ydoe):
+                f.write(f"{xd[0]:.6e} {xd[1]:.6e} {yd[0]:.6e}\n")
+    else:
+        print("Loading existing DOE...")
+        initial_dataset = np.load(f'{surrogate_folder_name}/ego_doe.npy')
+        xdoe = initial_dataset[:, :xlimits.shape[0]]
+        ydoe = initial_dataset[:, xlimits.shape[0]:]
+
+    initial_dataset = np.hstack((xdoe,ydoe))
+    #save the doe
+    np.save(f'{surrogate_folder_name}/ego_doe.npy', initial_dataset)
+
+    trained_surrogate_model = leave_one_out(initial_dataset)
+
+    print("Starting EGO optimization...")
+    x_opt, y_opt, ind_best, x_data, y_data = run_ego(surrogate_model=trained_surrogate_model, 
+                                                    n_ego_iter=n_ego_iter, 
+                                                    xlimits=xlimits, 
+                                                    simulation=simulation)
@@ -34,11 +34,11 @@ wavelength = 1.55
 #define the sources
 source1 = pid.Source(input_field='fiber_modes/LP11_field_distribution_lda_1550_N_1025.npy', 
                     wavelength=wavelength, 
-                     space_grid=space_grid, 
+                     #space_grid=space_grid, 
                     field_norm=True)
 source2 = pid.Source(input_field='fiber_modes/LP21_field_distribution_lda_1550_N_1025.npy', 
                     wavelength=wavelength, 
-                     space_grid=space_grid, 
+                     #space_grid=space_grid, 
                     field_norm=True)


@@ -46,20 +46,20 @@ target = pid.Target(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_
                    center=np.array([20,0]),
                    associated_source=source1, 
                    wavelength=wavelength, 
-                    space_grid=space_grid, 
+                    #space_grid=space_grid, 
                    field_norm=True)
 target2 = pid.Target(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy', 
                     center=np.array([-20,0]),
                     associated_source=source2, 
                     wavelength=wavelength, 
-                     space_grid=space_grid, 
+                     #space_grid=space_grid, 
                     field_norm=True) 

 #make a plot of sources and targets
-target.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/target_plot")
-target2.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/target2_plot")
-source1.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/source1_plot")
-source2.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/source2_plot")
+#target.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/target_plot")
+#target2.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/target2_plot")
+#source1.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/source1_plot")
+#source2.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/source2_plot")

 sim = pid.Simulation(
    sources=[source1, source2],
@@ -78,7 +78,7 @@ sim.set_optimization(opt_param_min=opt_param_min,
                     phase_mask_init_size=output_fiber_core_diameter,
                     opt_max_eval=opt_max_eval,
                     strategy = 'mean',
-                     conic_filter_radius = 1.0, #i.e. 700 [nm]
+                     #conic_filter_radius = 1.0, #i.e. 700 [nm]
                    )

 sim.run()