diff --git a/mpid.py b/mpid.py index 3bfc951..d1f7025 100644 --- a/mpid.py +++ b/mpid.py @@ -75,27 +75,6 @@ class SpaceGrid: Spatial dimensions are intended to be measured in micrometers [um]. ''' 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 - self.nx = (self.Nx - 1) // 2 - self.ny = (self.Ny - 1) // 2 - self.xy_symmetry = xy_symmetry - self.S = self.Nx * self.Ny - - if self.xy_symmetry: - self.dof = 2*((self.nx + 1)**2) - else: - self.dof = 2 * self.S - - 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 - if prop_distances is None: self.propagation_distances_set = False self.d = None @@ -105,6 +84,40 @@ class SpaceGrid: self.materials_refr_index = materials_refr_index + self.num_phase_masks = 0 + if self.propagation_distances_set: + self.num_phase_masks = len(self.d) - 1 + + #determine the number of phase masks required based on the number of propagation distances + self.num_phase_masks = 0 + if self.propagation_distances_set: + self.num_phase_masks = len(self.d) - 1 + + 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 # convert to [m] + self.nx = (self.Nx - 1) // 2 + self.ny = (self.Ny - 1) // 2 + self.xy_symmetry = xy_symmetry + self.S = self.Nx * self.Ny + + if self.xy_symmetry: + self.dof = self.num_phase_masks * ((self.nx + 1)**2) + self.S_symm = (self.nx + 1)**2 + else: + self.dof = self.num_phase_masks * self.S + + 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 + + def get_num_phase_masks(self): + return self.num_phase_masks + def get_grid_shape(self): return (self.Nx, self.Ny) @@ -121,10 +134,16 @@ class SpaceGrid: return self.sampling_period def get_resolution(self): - return self.unit_cell_pitch + ''' + Returns the resolution of the spatial grid in micrometers [um]. + ''' + return self.unit_cell_pitch * 1e6 # convert to [um] - def get_unit_cell_pitch(self): - return self.unit_cell_pitch + def get_unit_cell_pitch(self) -> float: + ''' + Returns the unit cell pitch of the spatial grid in micrometers [um]. + ''' + return self.unit_cell_pitch * 1e6 # convert back to [um] def get_size_x(self): return self.size_x @@ -150,6 +169,11 @@ class SpaceGrid: def get_S(self): return self.S + def get_S_symm(self): + if not self.xy_symmetry: + raise ValueError("S_symm is only defined when xy_symmetry is set True") + return self.S_symm + def get_X(self): return self.X @@ -320,7 +344,7 @@ class Simulation: if space_grid is not None: self.set_space_grid(space_grid) - + self._configure_fftw(self.n_cpu) self.is_optimization_set = False @@ -429,6 +453,9 @@ class Simulation: def _phase_given_w(self, input_normalized_phase): return 2 * np.pi * input_normalized_phase + def _w_given_phase(self, input_phase): + return input_phase / (2 * np.pi) + 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. @@ -467,7 +494,7 @@ class Simulation: 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())) + self.w_norm = np.zeros((self.num_phase_masks*self.space_grid.get_S())) circular_mask_source = [self._build_circular_mask(source_center, self.phase_mask_init_size) for source_center in source_centers] circular_mask_target = [self._build_circular_mask(target_center, self.phase_mask_init_size) for target_center in target_centers] @@ -483,17 +510,29 @@ class Simulation: self.w_norm[:self.space_grid.get_S()] = ((cumulated_source_phase + 2*np.pi) % (2*np.pi)) / (2*np.pi) #initialize mask on target side - cumulated_target_phase = np.zeros_like(self.w_norm[self.space_grid.get_S():]) + cumulated_target_phase = np.zeros(self.space_grid.get_S()) for _, target_phase in enumerate(target_pattern_phase): cumulated_target_phase += target_phase / (2 * np.pi) #normalize the target phase to be between 0 and 1 - self.w_norm[self.space_grid.get_S():] = ((cumulated_target_phase + 2*np.pi) % (2*np.pi)) / (2*np.pi) + if self.num_phase_masks ==2: + self.w_norm[self.space_grid.get_S():] = ((cumulated_target_phase + 2*np.pi) % (2*np.pi)) / (2*np.pi) + elif self.num_phase_masks == 3: + self.w_norm[2*self.space_grid.get_S():] = ((cumulated_target_phase + 2*np.pi) % (2*np.pi)) / (2*np.pi) + if self.space_grid.xy_symmetry: - 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():]) + if self.num_phase_masks == 2: + 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():]) + elif self.num_phase_masks == 3: + 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():2*self.space_grid.get_S()]) + w_norm_symm_3 = self._symmetrize(self.w_norm[2*self.space_grid.get_S():]) - self.w_norm_symm[:] = np.concatenate((w_norm_symm_1, w_norm_symm_2), axis=0) + if self.num_phase_masks == 2: + self.w_norm_symm[:] = np.concatenate((w_norm_symm_1, w_norm_symm_2), axis=0) + elif self.num_phase_masks == 3: + self.w_norm_symm[:] = np.concatenate((w_norm_symm_1, w_norm_symm_2, w_norm_symm_3), axis=0) self.phase_mask[:] = self._phase_given_w(self.w_norm) @@ -529,10 +568,20 @@ class Simulation: ) return npa.arctan2(phase_mask_filtered_sin, phase_mask_filtered_cos) - def _forward_propagate(self, source_object: Source): + def _forward_propagate(self, source_object: Source, returns_field: bool = False): #source field source_to_propagate = source_object.get_field(space_grid=self.space_grid) + #we retrieve the phase masks by managing the filtered phase mask array + #depending on the number of phase masks, we will split the filtered phase mask array accordingly + if self.num_phase_masks == 2: + filtered_phase_mask_1 = self.filtered_phase_mask[:self.space_grid.get_S()] + filtered_phase_mask_2 = self.filtered_phase_mask[self.space_grid.get_S():] + elif self.num_phase_masks == 3: + filtered_phase_mask_1 = self.filtered_phase_mask[:self.space_grid.get_S()] + filtered_phase_mask_2 = self.filtered_phase_mask[self.space_grid.get_S():2*self.space_grid.get_S()] + filtered_phase_mask_3 = self.filtered_phase_mask[2*self.space_grid.get_S():] + self.fft_input_array[:,:] = np.fft.ifftshift(source_to_propagate.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))) self.fft_operator() @@ -545,7 +594,7 @@ class Simulation: 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_input_array[:,:] = np.fft.ifftshift((input_field * np.exp(1j * filtered_phase_mask_1)).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() @@ -555,16 +604,35 @@ class Simulation: #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_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * np.exp(1j * filtered_phase_mask_2.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 + #output field if the number of phase masks is 2, otherwise this is the input field on MS3 self.intermediate_electric_fields[2] = np.fft.fftshift(self.ifft_operator.output_array.copy()) - self.output_field = self.intermediate_electric_fields[-1].flatten() + #if the number of phase masks is 2, stop here + if self.num_phase_masks == 2: + self.output_field = self.intermediate_electric_fields[-1].flatten() + if returns_field: + return self.output_field + else: + return None + self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * np.exp(1j * filtered_phase_mask_3.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[3] #element wise product + self.ifft_operator() + + self.intermediate_electric_fields[3] = np.fft.fftshift(self.ifft_operator.output_array.copy()) + + self.output_field = self.intermediate_electric_fields[-1].flatten() + + if returns_field: + return self.output_field + else: + return None #make a plot of 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") @@ -589,25 +657,58 @@ class Simulation: 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.filtered_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.filtered_phase_mask[self.space_grid.get_S():]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + #we need to manage the case of two or three phase masks by defining the associated filtered phase masks + if self.num_phase_masks == 2: + filtered_phase_mask_1 = self.filtered_phase_mask[:self.space_grid.get_S()] + filtered_phase_mask_2 = self.filtered_phase_mask[self.space_grid.get_S():] + elif self.num_phase_masks == 3: + filtered_phase_mask_1 = self.filtered_phase_mask[:self.space_grid.get_S()] + filtered_phase_mask_2 = self.filtered_phase_mask[self.space_grid.get_S():2*self.space_grid.get_S()] + filtered_phase_mask_3 = self.filtered_phase_mask[2*self.space_grid.get_S():] + Phi_dagger = [] + Phi_1_dagger = np.exp(-1j * filtered_phase_mask_1).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + Phi_2_dagger = np.exp(-1j * filtered_phase_mask_2).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + Phi_dagger.append(Phi_1_dagger) + Phi_dagger.append(Phi_2_dagger) + if self.num_phase_masks == 3: + Phi_3_dagger = np.exp(-1j * filtered_phase_mask_3).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + Phi_dagger.append(Phi_3_dagger) + + + grad_C_phi = [None] * self.num_phase_masks #compute the adjoint source - 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_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_electric_fields[0].conj() * np.fft.fftshift(self.ifft_operator.output_array) * Phi_1_dagger).flatten() + S = np.sum(self.intermediate_electric_fields[-1] * np.conj(target.get_field(space_grid=self.space_grid))) + a = S * target.get_field(space_grid=self.space_grid) - return grad_C_phi_1, grad_C_phi_2 + self.fft_input_array[:,:] = np.fft.ifftshift(a) + self.fft_operator() + self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[-1] + self.ifft_operator() + self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * Phi_dagger[-1]) + + grad_C_phi[-1] = 2 * np.real(-1j * self.intermediate_electric_fields[-2].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[-2] #element wise + self.ifft_operator() + self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * Phi_dagger[-2]) + + grad_C_phi[-2] = 2 * np.real(-1j * self.intermediate_electric_fields[-3].conj() * np.fft.fftshift(self.fft_input_array)).flatten() + + if self.num_phase_masks == 2: + return grad_C_phi + + self.fft_operator() + self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[-3] #element wise + self.ifft_operator() + self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * Phi_dagger[-3]) + + grad_C_phi[-3] = 2 * np.real(-1j * self.intermediate_electric_fields[-4].conj() * np.fft.fftshift(self.fft_input_array)).flatten() + + return grad_C_phi + #return grad_C_phi_1, grad_C_phi_2 def _value_and_grad(self, target): ''' @@ -621,9 +722,9 @@ class Simulation: self._forward_propagate(source) 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) + grads = self._adjoint_propagate(target) - return C_s, grad_C_phi_1, grad_C_phi_2 + return C_s, grads def _expand_symmetrize(self, w_symm): ''' @@ -661,28 +762,52 @@ class Simulation: print("evaluating epigraph constraint") epigraph = epigraph_and_weights[0] x = epigraph_and_weights[1:] - + #we need to manage the number of phase masks, depending on the current object instance if self.space_grid.xy_symmetry: - 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) + if self.num_phase_masks == 2: + ms1_first_quadrant = x[:self.space_grid.get_S_symm()].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + ms2_first_quadrant = x[self.space_grid.get_S_symm():].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) + elif self.num_phase_masks == 3: + ms1_first_quadrant = x[:self.space_grid.get_S_symm()].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + ms2_first_quadrant = x[self.space_grid.get_S_symm(): 2*self.space_grid.get_S_symm()].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + ms3_first_quadrant = x[2*self.space_grid.get_S_symm():].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) + ms3 = self._expand_symmetrize(ms3_first_quadrant) else: - ms1 = x[:self.space_grid.get_dof()//2] - ms2 = x[self.space_grid.get_dof()//2:] + if self.num_phase_masks == 2: + ms1 = x[:self.space_grid.get_S()] + ms2 = x[self.space_grid.get_S():] + elif self.num_phase_masks == 3: + ms1 = x[:self.space_grid.get_S()] + ms2 = x[self.space_grid.get_S(): 2*self.space_grid.get_S()] + ms3 = x[2*self.space_grid.get_S():] - all_ms = np.concatenate((ms1, ms2)) + if self.num_phase_masks == 2: + all_ms = np.concatenate((ms1, ms2)) + elif self.num_phase_masks == 3: + all_ms = np.concatenate((ms1, ms2, ms3)) self.w_norm[:] = all_ms 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)) + if self.num_phase_masks == 2: + 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)) + elif self.num_phase_masks == 3: + phase_mask_1 = self.phase_mask[:self.space_grid.get_S()] + phase_mask_2 = self.phase_mask[self.space_grid.get_S(): 2*self.space_grid.get_S()] + phase_mask_3 = self.phase_mask[2*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) + filtered_phase_mask_3 = self._phase_filter(phase_mask_3) + self.filtered_phase_mask[:] = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2, filtered_phase_mask_3)) else: self.filtered_phase_mask[:] = self.phase_mask @@ -692,25 +817,45 @@ class Simulation: for idx, target in enumerate(self.targets): - C, full_grad_1, full_grad_2 = self._value_and_grad(target) + C, full_grads = self._value_and_grad(target) #minus sign required because the constraint is in the form of f - t < 0 and we minimize the epigraph t results_array[idx] = - C if self.use_filtering: - jac_grad_1 = tensor_jacobian_product(self._phase_filter)(phase_mask_1, full_grad_1) - print("jac 1 done") - jac_grad_2 = tensor_jacobian_product(self._phase_filter)(phase_mask_2, full_grad_2) - print("jac 2 done") + if self.num_phase_masks == 2: + jac_grad_1 = tensor_jacobian_product(self._phase_filter)(phase_mask_1, full_grads[0]) + print("jac 1 done") + jac_grad_2 = tensor_jacobian_product(self._phase_filter)(phase_mask_2, full_grads[1]) + print("jac 2 done") + elif self.num_phase_masks == 3: + jac_grad_1 = tensor_jacobian_product(self._phase_filter)(phase_mask_1, full_grads[0]) + print("jac 1 done") + jac_grad_2 = tensor_jacobian_product(self._phase_filter)(phase_mask_2, full_grads[1]) + print("jac 2 done") + jac_grad_3 = tensor_jacobian_product(self._phase_filter)(phase_mask_3, full_grads[2]) + print("jac 3 done") else: - jac_grad_1 = full_grad_1 - jac_grad_2 = full_grad_2 + if self.num_phase_masks == 2: + jac_grad_1 = full_grads[0] + jac_grad_2 = full_grads[1] + elif self.num_phase_masks == 3: + jac_grad_1 = full_grads[0] + jac_grad_2 = full_grads[1] + jac_grad_3 = full_grads[2] if self.space_grid.xy_symmetry: jac_grad_1_symm = self._reduce_symmetrized_gradient(jac_grad_1) jac_grad_2_symm = self._reduce_symmetrized_gradient(jac_grad_2) - jac_grad = np.concatenate((jac_grad_1_symm, jac_grad_2_symm), axis=0) + if self.num_phase_masks == 3: + jac_grad_3_symm = self._reduce_symmetrized_gradient(jac_grad_3) + jac_grad = np.concatenate((jac_grad_1_symm, jac_grad_2_symm, jac_grad_3_symm), axis=0) + else: + jac_grad = np.concatenate((jac_grad_1_symm, jac_grad_2_symm), axis=0) else: - jac_grad = np.concatenate((jac_grad_1, jac_grad_2),axis=0) + if self.num_phase_masks == 2: + jac_grad = np.concatenate((jac_grad_1, jac_grad_2), axis=0) + elif self.num_phase_masks == 3: + jac_grad = np.concatenate((jac_grad_1, jac_grad_2, jac_grad_3), axis=0) grads_matrix[idx, :] = jac_grad @@ -749,26 +894,53 @@ class Simulation: def _cost_fun(self, x, grad): self.iter_num[0] = self.iter_num[0] + 1 if self.space_grid.xy_symmetry: - 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) + if self.num_phase_masks == 2: + ms1_first_quadrant = x[:self.space_grid.get_S_symm()].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + ms2_first_quadrant = x[self.space_grid.get_S_symm():].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) + elif self.num_phase_masks == 3: + ms1_first_quadrant = x[:self.space_grid.get_S_symm()].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + ms2_first_quadrant = x[self.space_grid.get_S_symm():2*self.space_grid.get_S_symm()].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + ms3_first_quadrant = x[2*self.space_grid.get_S_symm():].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) + ms3 = self._expand_symmetrize(ms3_first_quadrant) else: - ms1 = x[:self.space_grid.get_dof()//2] - ms2 = x[self.space_grid.get_dof()//2:] + if self.num_phase_masks == 2: + ms1 = x[:self.space_grid.get_dof()//2] + ms2 = x[self.space_grid.get_dof()//2:] + elif self.num_phase_masks == 3: + ms1 = x[:self.space_grid.get_dof()//3] + ms2 = x[self.space_grid.get_dof()//3:2*self.space_grid.get_dof()//3] + ms3 = x[2*self.space_grid.get_dof()//3:] - all_ms = np.concatenate((ms1, ms2)) + if self.num_phase_masks == 2: + all_ms = np.concatenate((ms1, ms2)) + elif self.num_phase_masks == 3: + all_ms = np.concatenate((ms1, ms2, ms3)) self.w_norm[:] = all_ms 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) + if self.num_phase_masks == 2: + 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) + elif self.num_phase_masks == 3: + phase_mask_1 = self.phase_mask[:self.space_grid.get_S()] + phase_mask_2 = self.phase_mask[self.space_grid.get_S():2*self.space_grid.get_S()] + phase_mask_3 = self.phase_mask[2*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) + filtered_phase_mask_3 = self._phase_filter(phase_mask_3) - self.filtered_phase_mask[:] = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2)) + if self.num_phase_masks == 2: + self.filtered_phase_mask[:] = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2)) + elif self.num_phase_masks == 3: + self.filtered_phase_mask[:] = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2, filtered_phase_mask_3)) else: self.filtered_phase_mask[:] = self.phase_mask @@ -779,23 +951,44 @@ class Simulation: for target in self.targets: - C, full_grad_1, full_grad_2 = self._value_and_grad(target) + C, full_grads = 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) + if self.num_phase_masks == 2: + full_grad_1, full_grad_2 = full_grads + 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) + elif self.num_phase_masks == 3: + full_grad_1, full_grad_2, full_grad_3 = full_grads + 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) + jac_grad_3 = tensor_jacobian_product(self._phase_filter)(phase_mask_3, full_grad_3) else: - jac_grad_1 = full_grad_1 - jac_grad_2 = full_grad_2 + if self.num_phase_masks == 2: + full_grad_1, full_grad_2 = full_grads + jac_grad_1 = full_grad_1 + jac_grad_2 = full_grad_2 + elif self.num_phase_masks == 3: + full_grad_1, full_grad_2, full_grad_3 = full_grads + jac_grad_1 = full_grad_1 + jac_grad_2 = full_grad_2 + jac_grad_3 = full_grad_3 if self.space_grid.xy_symmetry: jac_grad_1_symm = self._reduce_symmetrized_gradient(jac_grad_1) jac_grad_2_symm = self._reduce_symmetrized_gradient(jac_grad_2) - jac_grad = np.concatenate((jac_grad_1_symm, jac_grad_2_symm), axis=0) + if self.num_phase_masks == 2: + jac_grad = np.concatenate((jac_grad_1_symm, jac_grad_2_symm), axis=0) + elif self.num_phase_masks == 3: + jac_grad_3_symm = self._reduce_symmetrized_gradient(jac_grad_3) + jac_grad = np.concatenate((jac_grad_1_symm, jac_grad_2_symm, jac_grad_3_symm), axis=0) else: - jac_grad = np.concatenate((jac_grad_1, jac_grad_2),axis=0) + if self.num_phase_masks == 2: + jac_grad = np.concatenate((jac_grad_1, jac_grad_2),axis=0) + elif self.num_phase_masks == 3: + jac_grad = np.concatenate((jac_grad_1, jac_grad_2, jac_grad_3),axis=0) grad_partial[:] += jac_grad @@ -840,19 +1033,25 @@ class Simulation: ''' 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)) + if self.num_phase_masks == 2: + 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)) + elif self.num_phase_masks == 3: + 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():2*self.space_grid.get_S()]) + filtered_phase_mask_3 = self._phase_filter(self.phase_mask[2*self.space_grid.get_S():]) + self.filtered_phase_mask = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2, filtered_phase_mask_3)) 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_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.tile(mask_upper_bounds,self.num_phase_masks) 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_upper_bounds_symm_full = np.tile(mask_upper_bounds_symm,self.num_phase_masks) 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 @@ -923,23 +1122,45 @@ class Simulation: 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="phase", save_name=f"{self.results_folder}/output_fields_{i}_phase_plot") #phase masks - 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())) + if self.num_phase_masks == 2: + 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())) + elif self.num_phase_masks == 3: + 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():2*self.space_grid.get_S()]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + phase_mask_3 = self._phase_given_w(self.w_norm[2*self.space_grid.get_S():]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + if self.use_filtering: - filtered_phase_mask_1 = self._phase_filter(phase_mask_1) - filtered_phase_mask_2 = self._phase_filter(phase_mask_2) + if self.num_phase_masks == 2: + filtered_phase_mask_1 = self._phase_filter(phase_mask_1) + filtered_phase_mask_2 = self._phase_filter(phase_mask_2) + elif self.num_phase_masks == 3: + filtered_phase_mask_1 = self._phase_filter(phase_mask_1) + filtered_phase_mask_2 = self._phase_filter(phase_mask_2) + filtered_phase_mask_3 = self._phase_filter(phase_mask_3) else: - filtered_phase_mask_1 = phase_mask_1 - filtered_phase_mask_2 = phase_mask_2 + if self.num_phase_masks == 2: + filtered_phase_mask_1 = phase_mask_1 + filtered_phase_mask_2 = phase_mask_2 + elif self.num_phase_masks == 3: + filtered_phase_mask_1 = phase_mask_1 + filtered_phase_mask_2 = phase_mask_2 + filtered_phase_mask_3 = phase_mask_3 make_2Dplot_of(phase_mask_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="phase_mask", save_name=f"{self.results_folder}/phase_mask_1_plot") make_2Dplot_of(phase_mask_2, 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_mask", save_name=f"{self.results_folder}/phase_mask_2_plot") make_2Dplot_of(filtered_phase_mask_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="phase_mask", save_name=f"{self.results_folder}/filtered_phase_mask_1_plot") make_2Dplot_of(filtered_phase_mask_2, 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_mask", save_name=f"{self.results_folder}/filtered_phase_mask_2_plot") + if self.num_phase_masks == 3: + make_2Dplot_of(phase_mask_3, 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_mask", save_name=f"{self.results_folder}/phase_mask_3_plot") + make_2Dplot_of(filtered_phase_mask_3, 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_mask", save_name=f"{self.results_folder}/filtered_phase_mask_3_plot") 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) np.save(f"{self.results_folder}/filtered_phase_mask_1.npy", filtered_phase_mask_1) np.save(f"{self.results_folder}/filtered_phase_mask_2.npy", filtered_phase_mask_2) + if self.num_phase_masks == 3: + np.save(f"{self.results_folder}/optimized_phase_mask_3.npy", phase_mask_3) + np.save(f"{self.results_folder}/filtered_phase_mask_3.npy", filtered_phase_mask_3) def get_best_opt_val(self): if self.is_optimized_flag: @@ -968,6 +1189,7 @@ class Simulation: raise ValueError("The provided space grid is not compatible with the target fields.") self.space_grid = new_space_grid + self.num_phase_masks = self.space_grid.get_num_phase_masks() self.space_grid_is_set = False self.is_optimized_flag = False @@ -989,12 +1211,14 @@ class Simulation: 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 = np.zeros((self.num_phase_masks*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] + self.intermediate_electric_fields = [None] * (self.num_phase_masks + 1) + self.intermediate_magnetic_fields = [None] * (self.num_phase_masks + 1) + #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) @@ -1017,7 +1241,7 @@ class Simulation: 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, + unit_cell_pitch=self.space_grid.get_unit_cell_pitch(), prop_distances=new_distances, materials_refr_index=self.space_grid.get_materials_refr_index(), xy_symmetry=True, @@ -1025,6 +1249,67 @@ class Simulation: self.set_space_grid(new_space_grid) + def inject_phase_masks(self, input_phase_masks: list): + """ + Inject the specified phase masks into the simulation. + Args: + input_phase_masks (list): A list of phase masks, flattened, to be injected into the simulation. + Returns: + None + """ + #check if the length of the list of input_phase_masks we are trying to inject + #has the same number of phase masks in this simulation object. + if len(input_phase_masks) != self.num_phase_masks: + raise ValueError(f"Expected {self.num_phase_masks} phase masks, but got {len(input_phase_masks)}.") + for i, phase_mask in enumerate(input_phase_masks): + self.phase_mask[i*self.space_grid.get_S():(i+1)*self.space_grid.get_S()] = phase_mask + self.w_norm[i*self.space_grid.get_S():(i+1)*self.space_grid.get_S()] = self._w_given_phase(phase_mask) + if self.use_filtering: + if self.num_phase_masks == 2: + 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)) + elif self.num_phase_masks == 3: + 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():2*self.space_grid.get_S()]) + filtered_phase_mask_3 = self._phase_filter(self.phase_mask[2*self.space_grid.get_S():]) + self.filtered_phase_mask = np.concatenate((filtered_phase_mask_1, filtered_phase_mask_2, filtered_phase_mask_3)) + else: + self.filtered_phase_mask = self.phase_mask + + + def source_sweep(self, distances: np.ndarray = None, input_phase_masks: list = None): + """ + This method runs the propagation algorithm for the specified sources and targets in the simulation object. + For every propagation step, it enforces the same phase masks. It can optionally update propagation distances and phase masks. + If not specified, it is assumed the simulation object has already been optimized and that the required quantities are already set. + It is useful for studying the wavelength-dependent performance of the system. + Args: + distances (np.ndarray, optional): Array containing the propagation distances. If None, the current distances are used. + input_phase_masks (list, optional): A list of phase masks to be injected into the simulation. If None, the current phase masks are used. + Returns: + dict: A dictionary mapping source wavelengths to the overlap between the output field and the associated target field. + """ + results_dict = {} + + if distances is not None: + self.update_propagation_distances(distances) + + if input_phase_masks is not None: + self.inject_phase_masks(input_phase_masks) + + for target in self.targets: + source = target.get_associated_source() + self._update_wavelength(source.get_wavelength()) + output_field = self._forward_propagate(source, returns_field = True) + #compute the overlap between the output field and associated target field + target_field = target.get_field(self.get_space_grid(), flatten = True) + + overlap = np.abs(np.sum(np.conj(target_field) * output_field))**2 + results_dict[source.get_wavelength()] = overlap + + return results_dict + def run(self): """ Run the optimization process for the simulation. @@ -1033,6 +1318,7 @@ class Simulation: """ if self.space_grid_is_set: if self.is_optimization_set: + print(f"Using {self.num_phase_masks} phase masks") print("Initializing solver...") self._solver_init() print("Starting optimization...") @@ -1050,16 +1336,31 @@ class Simulation: 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.num_phase_masks == 2: + self.w_norm[:self.space_grid.get_S()] = self._expand_symmetrize(self.w_norm_symm[:self.space_grid.get_S_symm()]) + self.w_norm[self.space_grid.get_S():] = self._expand_symmetrize(self.w_norm_symm[self.space_grid.get_S_symm():]) + self.phase_mask[:] = self._phase_given_w(self.w_norm) + elif self.num_phase_masks == 3: + self.w_norm[:self.space_grid.get_S()] = self._expand_symmetrize(self.w_norm_symm[:self.space_grid.get_S_symm()]) + self.w_norm[self.space_grid.get_S():2*self.space_grid.get_S()] = self._expand_symmetrize(self.w_norm_symm[self.space_grid.get_S_symm():2*self.space_grid.get_S_symm()]) + self.w_norm[2*self.space_grid.get_S():] = self._expand_symmetrize(self.w_norm_symm[2*self.space_grid.get_S_symm():]) + 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]) + if self.num_phase_masks == 2: + 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]) + elif self.num_phase_masks == 3: + phase_mask_1 = self.phase_mask[:self.space_grid.get_S()] + phase_mask_2 = self.phase_mask[self.space_grid.get_S():2*self.space_grid.get_S()] + phase_mask_3 = self.phase_mask[2*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) + filtered_phase_mask_3 = self._phase_filter(phase_mask_3) + self.filtered_phase_mask[:] = np.concatenate([filtered_phase_mask_1, filtered_phase_mask_2, filtered_phase_mask_3]) else: self.filtered_phase_mask[:] = self.phase_mask @@ -1077,11 +1378,20 @@ class Simulation: 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]) + if self.num_phase_masks == 2: + 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]) + elif self.num_phase_masks == 3: + phase_mask_1 = self.phase_mask[:self.space_grid.get_S()] + phase_mask_2 = self.phase_mask[self.space_grid.get_S():2*self.space_grid.get_S()] + phase_mask_3 = self.phase_mask[2*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) + filtered_phase_mask_3 = self._phase_filter(phase_mask_3) + self.filtered_phase_mask[:] = np.concatenate([filtered_phase_mask_1, filtered_phase_mask_2, filtered_phase_mask_3]) else: self.filtered_phase_mask[:] = self.phase_mask diff --git a/single_converter_symm.py b/single_converter_symm.py index 1b2cc09..7d69ba1 100644 --- a/single_converter_symm.py +++ b/single_converter_symm.py @@ -7,16 +7,19 @@ if not os.path.exists("results"): N = 1025 unit_cell_pitch = 0.700 -n0 = 1.45 +n0 = 1.444 n1 = 1.0 -n2 = 1.45 +n2 = 1.444 mat_refr_index = np.array([n0, n1, n2]) #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 +#d0 = 516 #propagation distance: source - MS1 [um] FILTERED mf (using LP01 of FMF) +#d1 = 467 #propagation distance: MS1 - MS2 [um] FILTERED mf (using LP01 of FMF) +d0 = 195 #propagation distance: source - MS1 [um] FILTERED mf (using LP01 of SMF) +d1 = 129 #propagation distance: MS1 - MS2 [um] FILTERED mf (using LP01 of SMF) + d2 = d0 #propagation distance: MS2 - target [um] d = np.array([d0, d1, d2]) @@ -33,14 +36,10 @@ space_grid = pid.SpaceGrid( ) #define the source -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_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_files = ["fiber_modes/SMF/SMF_LP01_field_distribution_lda_1530_N_1025.npy", + "fiber_modes/SMF/SMF_LP01_field_distribution_lda_1550_N_1025.npy", + "fiber_modes/SMF/SMF_LP01_field_distribution_lda_1570_N_1025.npy"] sources = [pid.Source(input_field=np.load(file), wavelength=wvl, field_norm=True) for file, wvl in zip(source_files, wavelengths_to_test_um)] @@ -48,9 +47,9 @@ sources = [pid.Source(input_field=np.load(file), #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_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_files = ["fiber_modes/FMF/FMF_LP21_field_distribution_lda_1530_N_1025.npy", + "fiber_modes/FMF/FMF_LP21_field_distribution_lda_1550_N_1025.npy", + "fiber_modes/FMF/FMF_LP21_field_distribution_lda_1570_N_1025.npy"] target_fields = [np.load(file) for file in target_files] @@ -76,7 +75,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.0, + #conic_filter_radius=1.0, ) simulation.run() diff --git a/single_converter_symm_ego.py b/single_converter_symm_ego.py index 0c6830d..45d5d54 100644 --- a/single_converter_symm_ego.py +++ b/single_converter_symm_ego.py @@ -23,9 +23,9 @@ if not os.path.exists("results"): N = 1025 unit_cell_pitch = 0.700 #i.e., 700 [nm] -n0 = 1.45 +n0 = 1.444 n1 = 1.0 -n2 = 1.45 +n2 = 1.444 mat_refr_index = np.array([n0, n1, n2]) ms_max_available_diameter = 122 #[um] @@ -46,35 +46,25 @@ def setup_optical_system()->pid.Simulation: 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) + source_files = ["fiber_modes/SMF/SMF_LP01_field_distribution_lda_1530_N_1025.npy", + "fiber_modes/SMF/SMF_LP01_field_distribution_lda_1550_N_1025.npy", + "fiber_modes/SMF/SMF_LP01_field_distribution_lda_1570_N_1025.npy"] 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"] + target_files = ["fiber_modes/FMF/FMF_LP21_field_distribution_lda_1530_N_1025.npy", + "fiber_modes/FMF/FMF_LP21_field_distribution_lda_1550_N_1025.npy", + "fiber_modes/FMF/FMF_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, @@ -105,8 +95,8 @@ def setup_optical_system()->pid.Simulation: return simulation -def EI(GP, points, f_min): - pred = - GP.predict_values(points) #GP predicts overlaps (positive) +def EI(GP, points, f_min): #f_min is always negative + pred = - GP.predict_values(points) #GP predicts overlaps, which are always positive var = GP.predict_variances(points) args0 = (f_min - pred) / np.sqrt(var) args1 = (f_min - pred) * norm.cdf(args0) @@ -118,167 +108,134 @@ def EI(GP, points, f_min): 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) +def leave_one_out(dataset, corr_function): + full_x_doe = dataset[0] + full_y_doe = dataset[1] loo = LeaveOneOut() - #available kernels: ‘matern52’, ‘matern32’ ‘abs_exp’, ‘squar_exp’ - kernels_to_test = ['abs_exp', 'squar_exp', 'matern52', 'matern32'] - all_z_values = [] + z_values = [] + true_values = [] + predicted_values = [] + predicted_confidence = [] - 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=False, + corr=corr_function, + ) + sm_loo.set_training_values(train_values, train_outputs) + sm_loo.train() + #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}") - 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 + #SCATTER PLOT WITH ERROR BAR ON EACH SAMPLE loo cv 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) + 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.set_xticks(range(1, len(true_values), 2)) + ax.tick_params(axis='both', labelsize=20) ax.grid(True, alpha=0.3) ax.legend() - plt.savefig(f'{results_folder}/loo_standardized_residuals_gaussian_fit.pdf') + plt.savefig(f'{surrogate_folder_name}/loo_validation_kernel_{corr_function}.pdf') plt.close(fig) - best_kernel = kernels_to_test[np.argmin(np.array(all_std_values))] + #DIAGONAL COMPARISON LINE true/predicted + true_values_array = np.concatenate(true_values).ravel() + predicted_values_array = np.concatenate(predicted_values).ravel() - #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() + lower = min(true_values_array.min(), predicted_values_array.min()) + upper = max(true_values_array.max(), predicted_values_array.max()) - #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) + 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=20) + ax.grid(True, alpha=0.3) + ax.legend() + plt.savefig(f'{surrogate_folder_name}/loo_predicted_vs_true_diagonal_kernel_{corr_function}.pdf') + plt.close(fig) - #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 + #GAUSSIAN FIT standard normal vs. standardized residuals + fig, ax = plt.subplots(figsize=(10, 6)) + z_values_array = np.concatenate(z_values) + mu, std = norm.fit(z_values_array) + x = np.linspace(-4, 4, 100) + p = norm.pdf(x, mu, std) + ax.plot(x, p, 'r-', label='Fitted Gaussian') + mu_norm, std_norm = 0, 1 + ax.plot(x, norm.pdf(x, mu_norm, std_norm), 'g--', label='Standard Normal') + ax.hist(z_values_array, bins=20, density=True, alpha=0.6, color='blue', label='Standardized residuals') + ax.set_xlabel('Standardized Residual', fontsize=20) + ax.set_ylabel('Density', fontsize=20) + ax.set_title('Gaussian Fit vs Standardized Residuals', fontsize=20) + ax.tick_params(axis='both', labelsize=20) + ax.grid(True, alpha=0.3) + ax.legend() + plt.savefig(f'{surrogate_folder_name}/loo_gaussian_fit_kernel_{corr_function}.pdf') + plt.close(fig) -def run_ego(surrogate_model, n_ego_iter, xlimits, simulation): + #STANDARDIZED RESIDUALS PLOT + 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(1, len(z_values_array), 2)) + ax.set_ylim(-np.max(np.abs(z_values_array)), np.max(np.abs(z_values_array))) + ax.tick_params(axis='both', labelsize=20) + ax.grid(True, alpha=0.3) + ax.legend() + plt.savefig(f'{surrogate_folder_name}/loo_standardized_residuals_kernel_{corr_function}.pdf') + plt.close(fig) + + return 0 + +def run_ego(initial_dataset, 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 + xdoe_train_loaded = copy.deepcopy(initial_dataset[:, :-1]) + ydoe_train_loaded = copy.deepcopy(initial_dataset[:, -1]) # ydoe_train_loaded > 0 + x_data = copy.deepcopy(initial_dataset[:, :-1]) + y_data = copy.deepcopy(initial_dataset[:, -1]) # y_data > 0 + n_start = 40 #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 + 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) @@ -302,7 +259,7 @@ def run_ego(surrogate_model, n_ego_iter, xlimits, simulation): 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 + y_et_k = simulation.get_best_opt_val() #evaluate the objective function at the new point; y_et_k < 0 x_et_k_tested = np.atleast_2d(x_et_k) x_et_k_tested_history.append(x_et_k_tested) @@ -333,31 +290,52 @@ def run_ego(surrogate_model, n_ego_iter, xlimits, simulation): 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') + fig.savefig(f'{surrogate_folder_name}/optimisation_{k}.pdf') plt.close(fig) + fig, ax = plt.subplots(figsize=(8,6)) + d0_grid, d1_grid = np.meshgrid(np.linspace(xlimits[0][0], xlimits[0][1], 100), + np.linspace(xlimits[1][0], xlimits[1][1], 100)) + y_pred_grid = np.zeros_like(d0_grid) + for i in range(d0_grid.shape[0]): + for j in range(d0_grid.shape[1]): + y_pred_grid[i,j] = sm.predict_values(np.atleast_2d([d0_grid[i,j], d1_grid[i,j]])) + contour = ax.contourf(d0_grid, d1_grid, y_pred_grid, levels=50, cmap='viridis') + fig.colorbar(contour, ax=ax, label='predicted overlap') + 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.set_xlabel('d0 [um]', fontsize=15) + ax.set_ylabel('d1 [um]', fontsize=15) + ax.set_xlim(xlimits[0]) + ax.set_ylim(xlimits[1]) + ax.legend(loc='upper left', framealpha=0.5, fontsize=15) + for spine in ax.spines.values(): + spine.set_linewidth(1.5) + ax.set_aspect('equal') + fig.tight_layout() + fig.savefig(f'{surrogate_folder_name}/surrogate_landscape_{k}.pdf') + 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. + sm.set_training_values(x_data, y_data) #train the surrogate with positive values. + sm.train() #train the model with the new data point + print(f"{k+1}-th EGO refinement complete") - ind_best = np.argmin(y_data) + ind_best = np.argmax(y_data) x_opt = x_data[ind_best] - y_opt = - y_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") + image_pt = mpimg.imread(f"{surrogate_folder_name}/optimisation_{k}.pdf") im = plt.imshow(image_pt) ims.append([im]) @@ -366,7 +344,6 @@ def run_ego(surrogate_model, n_ego_iter, xlimits, simulation): 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') @@ -387,7 +364,7 @@ def run_ego(surrogate_model, n_ego_iter, xlimits, simulation): 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_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]') @@ -395,6 +372,7 @@ def run_ego(surrogate_model, n_ego_iter, xlimits, simulation): 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') @@ -412,33 +390,26 @@ def run_ego(surrogate_model, n_ego_iter, xlimits, simulation): 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')) + pickle.dump((x_data, y_data), open(f'{surrogate_folder_name}/refined_dataset.pkl', 'wb')) return x_opt, y_opt, ind_best, x_data, y_data +def train_surrogate_model(initial_dataset, corr_function): + full_x_doe = initial_dataset[:, :2] + full_y_doe = initial_dataset[:, 2].reshape(-1, 1) -############################## -#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 + sm = KRG(design_space = design_space, + n_start = 40, + print_global = False, + eval_noise = False, + corr = corr_function, + ) + sm.set_training_values(full_x_doe, full_y_doe) + sm.train() + + return sm - #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'): +def get_doe(): + if not os.path.exists(f'{doe_folder_name}/ego_doe.npy'): #build DOE print("Building DOE...") sampling = LHS(xlimits=xlimits, seed=seed) @@ -462,18 +433,57 @@ if __name__ == "__main__": 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') + initial_dataset = np.load(f'{doe_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) + np.save(f'{doe_folder_name}/ego_doe.npy', initial_dataset) + np.savetxt(f'{doe_folder_name}/ego_doe.txt', initial_dataset) + + return initial_dataset - trained_surrogate_model = leave_one_out(initial_dataset) +############################## +#EGO setup +############################## +if __name__ == "__main__": + + xlimits = np.array([[100,700], [100,700]]) # d0, d1 in [um] + seed = 777 + np.random.seed(seed) + design_space = DesignSpace(xlimits, seed=seed) + + ndoe = 20 #number of DOE samples + n_ego_iter = 15 #number of EGO refinement iterations + + corr_function_name = 'matern32' + doe_folder_name = 'DOE' + surrogate_folder_name = corr_function_name + '_surrogate_filtered' + if not os.path.exists(doe_folder_name): + os.makedirs(doe_folder_name) + if not os.path.exists(surrogate_folder_name): + os.makedirs(surrogate_folder_name) + + simulation = setup_optical_system() + + #DEFINE AN INITIAL DESIGN OF EXPERIMENT (DOE) + initial_dataset = get_doe() + + #TRAIN THE SURROGATE ON THE DOE using specified correlation function + trained_surrogate_model = train_surrogate_model(initial_dataset, + corr_function=corr_function_name) + + #RUN EGO REFINEMENT for n_ego_iter iterations print("Starting EGO optimization...") - x_opt, y_opt, ind_best, x_data, y_data = run_ego(surrogate_model=trained_surrogate_model, + x_opt, y_opt, ind_best, x_data, y_data = run_ego( + initial_dataset=initial_dataset, + surrogate_model=trained_surrogate_model, n_ego_iter=n_ego_iter, xlimits=xlimits, simulation=simulation) + + #EVALUATE THE MODEL PERFORMANCE BY RUNNING A LOO CROSS-VALIDATION + refined_dataset = pickle.load(open(f'{surrogate_folder_name}/refined_dataset.pkl', 'rb')) + leave_one_out(refined_dataset, corr_function_name) \ No newline at end of file diff --git a/single_converter_symm_source_sweep.py b/single_converter_symm_source_sweep.py new file mode 100644 index 0000000..16639ea --- /dev/null +++ b/single_converter_symm_source_sweep.py @@ -0,0 +1,108 @@ +import mpid as pid +import numpy as np +import matplotlib.pyplot as plt +import os +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.444 +n1 = 1.0 +n2 = 1.444 +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 + + wavelengths_to_test_um = [1.530, 1.550, 1.570] #[um] + + source_files = ["fiber_modes/SMF/SMF_LP01_field_distribution_lda_1530_N_1025.npy", + "fiber_modes/SMF/SMF_LP01_field_distribution_lda_1550_N_1025.npy", + "fiber_modes/SMF/SMF_LP01_field_distribution_lda_1570_N_1025.npy"] + + 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/FMF/FMF_LP21_field_distribution_lda_1530_N_1025.npy", + "fiber_modes/FMF/FMF_LP21_field_distribution_lda_1550_N_1025.npy", + "fiber_modes/FMF/FMF_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))] + + 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 + + +if __name__ == "__main__": + simulation = setup_optical_system() + + results_folder_name = 'results_nf' + + phase_mask_1 = np.load(f"{results_folder_name}/filtered_phase_mask_1.npy").flatten() + phase_mask_2 = np.load(f"{results_folder_name}/filtered_phase_mask_2.npy").flatten() + + d0 = 195 #propagation distance: source - MS1 [um] FILTERED mf (using LP01 of SMF) + d1 = 129 #propagation distance: MS1 - MS2 [um] FILTERED mf (using LP01 of SMF) + + d2 = d0 #propagation distance: MS2 - target [um] + d = np.array([d0, d1, d2]) + results = simulation.source_sweep(distances = d, input_phase_masks=[phase_mask_1, phase_mask_2]) + #make a plot of the results, Use the keys as x axis and values as y axis + fig, ax = plt.subplots(figsize=(6, 6)) + ax.plot(np.array(list(results.keys()))*1e9, np.array(list(results.values())), marker='o', linewidth=2) + #make the y tick labels to only have two decimal places + ax.yaxis.set_major_formatter(plt.FormatStrFormatter('%.3f')) + ax.set_xlabel('wavelength [nm]', fontsize=15) + ax.set_ylabel('overlap', fontsize=15) + ax.tick_params(axis='both', which='major', labelsize=15) + ax.set_aspect('auto') + plt.tight_layout() + plt.savefig(f"{results_folder_name}/source_sweep_results.pdf") \ No newline at end of file diff --git a/sorter_to_gauss.py b/sorter_to_gauss.py index 5741b48..07bafbc 100644 --- a/sorter_to_gauss.py +++ b/sorter_to_gauss.py @@ -1,26 +1,28 @@ import mpid as pid import numpy as np import os -if not os.path.exists("results"): - os.makedirs("results") +if not os.path.exists("results_sorter"): + os.makedirs("results_sorter") #Space grid definition N = 1025 unit_cell_pitch = 0.700 #i.e., 700 [nm] n0 = 1.45 -n1 = 1.0 +n1 = 1.45 n2 = 1.45 -mat_refr_index = np.array([n0, n1, n2]) +n3 = 1.45 +mat_refr_index = np.array([n0, n1, n2, n3]) -d0 = 476 #propagation distance: source - MS1 [um] -d1 = 510 #propagation distance: MS1 - MS2 [um] -d2 = d0 #propagation distance: MS2 - target [um] -d = np.array([d0, d1, d2]) +d0 = 500 #propagation distance: source - MS1 [um] +d1 = d0 #propagation distance: MS1 - MS2 [um] +d2 = d0 #propagation distance: MS2 - MS3 [um] +d3 = d0 #propagation distance: MS3 - target [um] +d = np.array([d0, d1, d2, d3]) -ms_max_available_diameter = 122 #[um] +ms_max_available_diameter = 250 #[um] output_fiber_core_diameter = 19 #[um] -opt_max_eval = 140 +opt_max_eval = 200 space_grid = pid.SpaceGrid( side_pixels=N, @@ -32,38 +34,44 @@ space_grid = pid.SpaceGrid( wavelength = 1.55 #define the sources -source1 = pid.Source(input_field='fiber_modes/LP11_field_distribution_lda_1550_N_1025.npy', +source1 = pid.Source(input_field='fiber_modes/FMF/FMF_LP01_field_distribution_lda_1550_N_1025.npy', wavelength=wavelength, #space_grid=space_grid, field_norm=True) -source2 = pid.Source(input_field='fiber_modes/LP21_field_distribution_lda_1550_N_1025.npy', +source2 = pid.Source(input_field='fiber_modes/FMF/FMF_LP11_field_distribution_lda_1550_N_1025.npy', + wavelength=wavelength, + #space_grid=space_grid, + field_norm=True) +source3 = pid.Source(input_field='fiber_modes/FMF/FMF_LP21_field_distribution_lda_1550_N_1025.npy', wavelength=wavelength, #space_grid=space_grid, field_norm=True) -target = pid.Target(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy', +target1 = pid.Target(input_field='fiber_modes/SMF/SMF_LP01_field_distribution_lda_1550_N_1025.npy', center=np.array([20,0]), associated_source=source1, wavelength=wavelength, #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]), +target2 = pid.Target(input_field='fiber_modes/SMF/SMF_LP01_field_distribution_lda_1550_N_1025.npy', + center=np.array([0,0]), associated_source=source2, wavelength=wavelength, #space_grid=space_grid, - field_norm=True) + field_norm=True) + +target3 = pid.Target(input_field='fiber_modes/SMF/SMF_LP01_field_distribution_lda_1550_N_1025.npy', + center=np.array([-20,0]), + associated_source=source3, + wavelength=wavelength, + #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") sim = pid.Simulation( - sources=[source1, source2], - targets=[target, target2], + sources=[source1, source2, source3], + targets=[target1, target2, target3], central_wavelength=wavelength, space_grid=space_grid, n_cpu=1 @@ -78,7 +86,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, ) sim.run()