first commit

mpid.py

@@ -0,0 +1,585 @@
+import numpy as np
+import matplotlib.pyplot as plt
+plt.rcParams.update({'font.size': 14})
+import os
+import pyfftw
+import nlopt
+
+def make_2Dplot_of(given_field,grid_X,grid_Y,plot_zoom_x,plot_zoom_y,choose_quantity="amplitude",save_name="plot") -> None:
+    '''Makes a plot of the given 2D field in a zoomed region. 
+    Saves the plot with the chosen filename to a PDF file. 
+    Args:
+        given_field: the 2D field to be plotted (flattened)
+        choose_quantity: either "amplitude" or "phase" to select quantity to plot
+        save_name: name of the file to save the plot (without extension)
+        plot_zoom_x: zoom level for the x-axis (in micrometers)
+        plot_zoom_y: zoom level for the y-axis (in micrometers)
+    Returns:
+        None'''
+    
+    given_field_2d = given_field.reshape((grid_X.shape[0], grid_X.shape[1])) #assure the field is 2D before plotting
+    fig, ax = plt.subplots(figsize=(6, 6))
+    if choose_quantity == "amplitude":
+        ax.pcolormesh(grid_X*1e6,grid_Y*1e6,np.abs(given_field_2d),rasterized=True)
+        #add a colorbar for amplitude
+        plt.colorbar(ax.pcolormesh(grid_X*1e6,grid_Y*1e6,np.abs(given_field_2d),rasterized=True), ax=ax)
+    elif choose_quantity == "phase":
+        ax.pcolormesh(grid_X*1e6,grid_Y*1e6,(np.angle(given_field_2d)+2*np.pi)%(2*np.pi), cmap='hsv', vmin=0, vmax=2*np.pi, rasterized=True)
+    elif choose_quantity == "phase_mask":
+        ax.pcolormesh(grid_X*1e6,grid_Y*1e6,(given_field_2d+2*np.pi)%(2*np.pi), cmap='hsv', vmin=0, vmax=2*np.pi, rasterized=True)
+    else:
+        raise ValueError("Unexpected value for choose_quantity")
+
+    ax.set_xlim(-plot_zoom_x, plot_zoom_x)
+    ax.set_ylim(-plot_zoom_y, plot_zoom_y)
+    ax.set_xlabel("x [um]")
+    ax.set_ylabel("y [um]")
+    plt.tight_layout()
+    plt.savefig(save_name + ".pdf")
+    plt.close()
+
+def make_1Dplot_of(x_axis,y_axis,plot_zoom_x,save_name="plot_1D",) -> None:
+    '''Makes a 1D plot of the given data and saves it to a PDF file.
+    Args:
+        x_axis: list of x-axis data (list of 1D arrays)
+        y_axis: list of y-axis data (list of 1D arrays)
+        plot_zoom_x: zoom level for the x-axis (in micrometers)
+        save_name: name of the file to save the plot (without extension)
+    Returns:
+        None
+    '''
+    if len(x_axis) != len(y_axis):
+        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])
+    plt.xlabel("x [um]")
+    plt.tight_layout()
+    plt.xlim(-plot_zoom_x, plot_zoom_x)
+    plt.savefig(save_name + ".pdf")
+    plt.close()
+
+class SpaceGrid:
+    def __init__(self, side_pixels: int, unit_cell_pitch: float, prop_distances: np.ndarray, materials_refr_index: np.ndarray):
+        self.Nx = side_pixels#for the moment we consider only square grids
+        self.Ny = side_pixels
+        self.unit_cell_pitch = unit_cell_pitch
+        self.nx = (self.Nx - 1) // 2
+        self.ny = (self.Ny - 1) // 2
+        self.dof = 2*((self.nx + 1)**2) #this needs to be computed depending on problem symmetry, which should be passed as argument (symmetry should be an instance of the Symmetry class)
+        self.S = self.Nx * self.Ny
+        self.size_x = self.unit_cell_pitch * self.Nx
+        self.size_y = self.unit_cell_pitch * self.Ny
+        self.xs = np.arange(-self.Nx//2+1, self.Nx//2+1) * (self.size_x / self.Nx)
+        self.ys = np.arange(-self.Ny//2+1, self.Ny//2+1) * (self.size_y / self.Ny)
+        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
+
+        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.")
+        self.materials_refr_index = materials_refr_index
+
+    def get_grid_shape(self):
+        return (self.Nx, self.Ny)
+
+    def get_materials_refr_index(self):
+        return self.materials_refr_index
+
+    def get_distances(self):
+        return self.d
+
+    def get_num_distances(self):
+        return len(self.d)
+    
+    def get_sampling_period(self):
+        return self.sampling_period
+
+    def get_Nx(self):
+        return self.Nx
+
+    def get_Ny(self):
+        return self.Ny
+
+    def get_nx(self):
+        return self.nx
+
+    def get_ny(self):
+        return self.ny
+
+    def get_dof(self):
+        return self.dof
+
+    def get_S(self):
+        return self.S
+
+    def get_X(self):
+        return self.X
+
+    def get_Y(self):
+        return self.Y
+
+    def get_rho(self):
+        return self.rho
+    
+class Field:
+    def __init__(self, input_field: np.ndarray, wavelength: float, space_grid: SpaceGrid):
+        self.input_field = input_field
+        self.wavelength = wavelength
+        self.space_grid = space_grid
+        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):
+        if flatten:
+            return self.input_field.flatten()
+        return self.input_field
+
+    def get_wavelength(self):
+        return self.wavelength
+
+    def plot2D(self, plot_zoom_x, plot_zoom_y, 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)
+
+class Source(Field):
+    def __init__(self, input_field: np.ndarray, wavelength: float, space_grid: SpaceGrid, wvl_unit_measure = 1,):
+        super().__init__(input_field, wavelength * wvl_unit_measure, space_grid)
+        self.type = "source"
+
+    def get_type(self):
+        return self.type
+
+class Target(Field):
+    def __init__(self, input_field: np.ndarray, associated_source: Source, wavelength: float, space_grid: SpaceGrid, wvl_unit_measure = 1):
+        super().__init__(input_field, wavelength * wvl_unit_measure, space_grid)
+        self.type = "target"
+        if associated_source.get_wavelength() != wavelength * wvl_unit_measure:
+            raise ValueError("Associated source must have the same wavelength as the target.")
+        else:
+            self.associated_source = associated_source
+        
+    def get_type(self):
+        return self.type
+
+    def get_associated_source(self):
+        return self.associated_source
+
+class PlotArray():
+    def __init__(self, input_field: np.ndarray, space_grid: SpaceGrid):
+        self.input_field = input_field
+        self.array_shape = input_field.shape
+        self.space_grid = space_grid
+        self.type = "plot_array"
+
+    def get_type(self):
+        return self.type
+
+    def plot2D(self, plot_zoom_x, plot_zoom_y, choose_quantity="amplitude", save_name="results/plot2D_out"):
+        if len(self.array_shape) == 2:
+            make_2Dplot_of(self.input_field, 
+                        self.space_grid.X, 
+                        self.space_grid.Y,
+                        plot_zoom_x, 
+                        plot_zoom_y, 
+                        choose_quantity, 
+                        save_name)
+        else:
+            raise ValueError("Input field must be 2D for plot2D.")
+
+    def plot1D(self, plot_zoom_x, save_name="results/plot1D_out"):
+        if len(self.array_shape) == 1:
+            make_1Dplot_of([self.space_grid.X.flatten()], 
+                           [self.input_field.flatten()], 
+                           plot_zoom_x, 
+                           save_name)
+        else:
+            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):
+        self.sources = sources
+        self.targets = targets
+        self.space_grid = space_grid
+        self.central_wavelength = central_wavelength #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
+        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
+        self.results_while_solving = True
+        self.opt_history = []
+        self.results_folder = "results"
+        if not os.path.exists(self.results_folder):
+            os.makedirs(self.results_folder)
+
+    def _configure_fftw(self, n_cpu = 1, planner_effort='FFTW_ESTIMATE'):
+        pyfftw.config.NUM_THREADS = n_cpu
+        pyfftw.interfaces.cache.disable()
+        pyfftw.config.PLANNER_EFFORT = planner_effort
+
+    def set_optimization(self, opt_param_min, opt_param_max, ms_max_avail_diameter, phase_mask_init_size, opt_max_eval, results_while_solving=True, custom_mask = None):
+        '''
+        Sets the optimization parameters for the simulation.
+        Args:
+            optimization_params (dict): A dictionary containing optimization parameters.
+        '''
+        self.is_optimization_set = True
+        #self.optimization_params = optimization_params
+        self.iter_num = [0]
+        self.norm_phase_min = opt_param_min
+        self.norm_phase_max = opt_param_max
+        self.ms_max_avail_diameter = ms_max_avail_diameter
+        self.phase_mask_init_size = phase_mask_init_size
+        self.opt_max_eval = opt_max_eval
+        self.results_while_solving = results_while_solving
+        self.custom_mask = custom_mask #future: one might be able to directly pass a custom 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_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()
+
+        self.mask_upper_bounds_symm = np.concatenate((mask_upper_bounds_symm, mask_upper_bounds_symm),axis=0)
+        self.mask_lower_bounds_symm = np.zeros_like(self.mask_upper_bounds_symm)
+        
+    def get_sources(self):
+        return self.sources
+
+    def get_targets(self):
+        return self.targets
+
+    def get_space_grid(self):
+        return self.space_grid
+
+    def get_central_wavelength(self):
+        return self.central_wavelength
+
+    def get_operating_wavelength(self):
+        return self.operating_wavelength
+
+    def get_k_vectors(self):
+        return self.k_vectors
+
+    def get_wvl_all_targets(self, unit_measure = 1):
+        return [target.get_wavelength() * unit_measure for target in self.targets]
+
+    def _update_wavelength(self, new_wavelength):
+        '''
+        Updates the operating wavelength and related vectors (lda_vector and k_vectors) based on the new wavelength.
+        The new_wavelength should be provided in meters.
+        Args:
+            new_wavelength (float): The new operating wavelength in meters.
+        '''
+        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.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))]
+
+    def _phase_given_w(self):
+        return 2 * np.pi * self.w_norm
+
+    def _phase_given_w(self, input_normalized_phase):
+        return 2 * np.pi * input_normalized_phase
+
+    def _symmetrize(self, w):
+        '''
+        Takes an input array, which must be full in size, turns it into a 2D array, and symmetrizes it along the x and y axes.
+        Returns the first quadrant of shape (nx,ny), including the zero axis, flattened
+        '''
+        try:
+            w_2d = w.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))
+        except:
+            raise ValueError(f"Wrong shape for input array. Current shape: {w.shape}")
+
+        w_symm = np.zeros((self.space_grid.get_nx()+1, self.space_grid.get_ny()+1))
+        w_symm[:,:] = w_2d[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1]
+
+        return w_symm.flatten()
+
+    def _initialize_phase_masks(self):
+        '''
+        Initializes the phase masks for the simulation.
+        This method should be called once, before iterating over the targets.
+        If multiple targets are present in this simulation instance, the first initialization uses the first target
+        '''
+        target_to_initialize = self.targets[0]
+        target_Efield = target_to_initialize.get_field(flatten=True)
+
+        self.w_norm = np.zeros((2*self.space_grid.get_S()))
+
+        circular_mask_1 = np.where(self.space_grid.get_rho().flatten() < self.phase_mask_init_size, 1, 0)
+        target_pattern_phase = (np.angle(target_Efield) + 2 * np.pi) % (2 * np.pi)*circular_mask_1 #make sure the phase is between 0 and 2pi
+
+        self.w_norm[self.space_grid.get_S():] += target_pattern_phase / (2 * np.pi) #normalize the target phase to be between 0 and 1
+
+        w_norm_symm_1 = self._symmetrize(self.w_norm[:self.space_grid.get_S()])
+        w_norm_symm_2 = self._symmetrize(self.w_norm[self.space_grid.get_S():])
+
+        self.w_norm_symm[:] = np.concatenate((w_norm_symm_1, w_norm_symm_2), axis=0)
+
+        self.phase_mask[:] = self._phase_given_w(self.w_norm)
+
+    def _forward_propagate(self, source_object: Source):
+
+        source_to_propagate = source_object.get_field()
+
+        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()
+        self.intermediate_fields[0] = np.fft.fftshift(self.ifft_operator.output_array)
+        input_field = self.intermediate_fields[0].flatten()
+
+        self.fft_input_array[:,:] = np.fft.ifftshift((input_field * np.exp(1j * self.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()
+        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.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * np.exp(1j * self.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()
+        self.intermediate_fields[2] = np.fft.fftshift(self.ifft_operator.output_array.copy())
+
+        self.output_field = self.intermediate_fields[-1].flatten()
+
+
+    def _reduce_symmetrized_gradient(self, g_full):
+        g_2d = g_full.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))
+
+        third = g_2d[self.space_grid.get_nx()+1:,:self.space_grid.get_ny()+1]
+        second = g_2d[:self.space_grid.get_nx()+1, self.space_grid.get_ny()+1:]
+        fourth = g_2d[self.space_grid.get_nx()+1:, self.space_grid.get_ny()+1:]
+
+        g_2d[:self.space_grid.get_nx(), self.space_grid.get_ny()+1:] += np.flip(fourth, axis = 0)
+        g_2d[:self.space_grid.get_nx(), :self.space_grid.get_ny()+1] += np.flip(third, axis = 0)
+        g_2d[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()] += np.flip(second, axis = 1)
+        
+        return g_2d[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1].flatten()
+
+    def _adjoint_propagate(self, target):
+        Phi_1_dagger = np.exp(-1j * self.phase_mask[:self.space_grid.get_S()]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))
+        Phi_2_dagger = np.exp(-1j * self.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_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() 
+
+        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()
+        
+        return grad_C_phi_1, grad_C_phi_2
+
+    def _value_and_grad(self, target):
+
+        source = target.get_associated_source()
+        wavelength = source.get_wavelength()
+
+        self._update_wavelength(wavelength)
+        self._forward_propagate(source)
+
+        C_s = np.real(np.abs(np.sum(self.output_field * np.conj(target.get_field(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
+
+    def _expand_symmetrize(self, w_symm):
+        '''
+        Takes a array of shape (nx+1, ny+1), representing the first quadrant, including zero axes, and 
+        expands it to a full 2D array of shape (Nx, Ny) by mirroring, flattened.
+        '''
+
+        w_symm_2d = w_symm.reshape((self.space_grid.get_nx()+1, self.space_grid.get_ny()+1))
+
+        w_full = np.zeros((self.space_grid.get_Nx(), self.space_grid.get_Ny()))
+
+        w_full[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1] = w_symm_2d
+        w_full[:self.space_grid.get_nx()+1, self.space_grid.get_ny()+1:] = np.flip(w_full[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()], axis=1)
+        w_full[self.space_grid.get_nx()+1:, :self.space_grid.get_ny()+1] = np.flip(w_full[:self.space_grid.get_nx(), :self.space_grid.get_ny()+1], axis=0)
+        w_full[self.space_grid.get_nx()+1:, self.space_grid.get_ny()+1:] = np.flip(np.flip(w_full[:self.space_grid.get_nx(), :self.space_grid.get_ny()], axis=1), axis=0)
+        w_full[self.space_grid.get_nx(), self.space_grid.get_ny()] = w_symm_2d[-1, -1]
+
+        return w_full.flatten()
+
+    def _cost_fun(self, x, grad):
+
+        ms1_first_quadrant = x[:self.space_grid.get_dof()//2].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1)
+        ms2_first_quadrant = x[self.space_grid.get_dof()//2:].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1)
+
+        ms1 = self._expand_symmetrize(ms1_first_quadrant)
+        ms2 = self._expand_symmetrize(ms2_first_quadrant)
+
+        all_ms = np.concatenate((ms1, ms2))
+
+        self.w_norm[:] = all_ms
+        self.phase_mask[:] = self._phase_given_w(self.w_norm)
+
+        #after the phase mask is set, we test the same mask across all targets
+        C_all_wavelengths = []
+        grad_partial = np.zeros_like(grad)
+
+        for target in self.targets:
+
+            C, full_grad1, full_grad2 = self._value_and_grad(target) 
+
+            C_all_wavelengths.append(C)
+
+            reduced_grad1 = self._reduce_symmetrized_gradient(full_grad1)
+            reduced_grad2 = self._reduce_symmetrized_gradient(full_grad2)
+            reduced_grads = np.concatenate((reduced_grad1, reduced_grad2), axis=0)
+            grad_partial[:] += reduced_grads * 2 * np.pi / len(self.targets) #average the gradient across all wavelengths
+
+        if grad.size > 0:
+            grad[:] = grad_partial
+
+        #the new cost function considers the average overlap across all wavelengths.
+        self.opt_history.append(np.mean(C_all_wavelengths)) 
+        self.iter_num[0] += 1
+        if self.iter_num[0] % 10 == 0 and self.results_while_solving:
+            fig, ax = plt.subplots()
+            ax.plot(self.opt_history)
+            ax.set_xlabel('Iteration')
+            ax.set_ylabel('Cost Function Value')
+            ax.set_title('Optimization History')
+            fig.tight_layout()
+            fig.savefig(f"{self.results_folder}/optimization_history.pdf")
+            plt.close(fig)
+
+            #make a plot of the cost function values across 
+            # all wavelengths for this iteration
+            fig, ax = plt.subplots()
+            ax.plot(self.get_wvl_all_targets(unit_measure=self.space_grid.m_to_nm), C_all_wavelengths, marker='o')
+            ax.set_xlabel(f'Wavelength [nm]')
+            ax.set_ylabel('Cost function value')
+            fig.tight_layout()
+            fig.savefig(f"{self.results_folder}/cost_function_across_wavelengths.pdf")
+            plt.close(fig)
+
+        return np.mean(C_all_wavelengths) 
+
+    def _export_results(self):
+        '''
+        Placeholder method for exporting the results of the optimization.
+        This method should be implemented to save the relevant data.
+        '''
+        phase_mask_1 = self._phase_given_w(self.w_norm[:self.space_grid.get_S()]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))
+        phase_mask_2 = self._phase_given_w(self.w_norm[self.space_grid.get_S():]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))
+        np.save(f"{self.results_folder}/optimized_phase_mask_1.npy", phase_mask_1)
+        np.save(f"{self.results_folder}/optimized_phase_mask_2.npy", phase_mask_2)
+
+    def get_best_opt_val(self):
+        if self.is_optimized_flag:
+            return self.best_opt_val
+        else:
+            raise RuntimeError("Simulation object has not been optimized")
+
+    def get_best_opt_and_distances(self):
+        if self.is_optimized_flag:
+            return self.best_opt_val, self.space_grid.d
+        else:
+            raise RuntimeError("Simulation object has not been optimized")
+
+    def is_optimized(self):
+        return self.is_optimized_flag
+
+    def update_space_grid(self, new_space_grid: SpaceGrid):
+        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)
+
+    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
+            solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof())
+            solver.set_lower_bounds(self.mask_lower_bounds_symm.flatten())
+            solver.set_upper_bounds(self.mask_upper_bounds_symm.flatten())
+            solver.set_max_objective(self._cost_fun)
+            solver.set_maxeval(self.opt_max_eval)
+            solver.set_param("dual_ftol_rel", 1e-7)
+            solver.set_param("verbosity",0)
+
+            self.w_norm_symm[:] = solver.optimize(self.w_norm_symm)
+            self.is_optimized_flag = True
+            self.best_opt_val = 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._export_results()
+        else:
+            raise RuntimeError("Optimization settings have not been specified.")
+        

testing.py

@@ -0,0 +1,78 @@
+import mpid as pid
+import numpy as np
+
+#Space grid definition
+N = 1025
+unit_cell_pitch = 700e-9
+
+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]
+d = np.array([d0, d1, d2])
+
+ms_max_available_diameter = 122e-6 #[um]
+output_fiber_core_diameter = 19e-6 #[um]
+opt_max_eval = 140
+
+space_grid = pid.SpaceGrid(
+    side_pixels=N, 
+    unit_cell_pitch = unit_cell_pitch, 
+    prop_distances=d, 
+    materials_refr_index=mat_refr_index,
+)
+
+#define the source
+beam_waist = 5.2e-6
+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, 
+                      wavelength=wvl, 
+                      space_grid=space_grid, 
+                      wvl_unit_measure=space_grid.nm) for wvl in wavelengths_to_test_nm]
+
+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_1025.npy"]
+
+target_fields = [np.load(file) for file in target_files]
+norm_target_fields = [field / np.sqrt(np.sum(np.abs(field)**2)) for field in target_fields]
+
+targets = [pid.Target(input_field=norm_target_fields[i], 
+                      associated_source=sources[i], 
+                      wavelength=wvl, 
+                      space_grid=space_grid, 
+                      wvl_unit_measure=space_grid.nm) for i, (file, wvl) in enumerate(zip(target_files, wavelengths_to_test_nm))]
+
+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,
+    space_grid=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,
+)
+
+simulation.run()
+
+best_opt_val, distances = simulation.get_best_opt_and_distances()
+print("Best optimization value:", best_opt_val)
+print("Distances:", distances)