From 1d843d2553bf9ab56623baab812533e5ffdfb185 Mon Sep 17 00:00:00 2001 From: Edoardo Vergerio Date: Mon, 7 Sep 2026 13:38:42 +0200 Subject: [PATCH] stable definition of mpid.py. Commit version before testing conic filtering and autograd ad --- mpid.py | 211 +++++++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 171 insertions(+), 40 deletions(-) diff --git a/mpid.py b/mpid.py index e866612..95f379c 100644 --- a/mpid.py +++ b/mpid.py @@ -5,7 +5,7 @@ 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: +def make_2Dplot_of(given_field:np.ndarray, grid_X:np.ndarray, grid_Y:np.ndarray, plot_zoom_x:float,plot_zoom_y:float,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: @@ -60,14 +60,20 @@ def make_1Dplot_of(x_axis,y_axis,plot_zoom_x,save_name="plot_1D",) -> None: plt.close() class SpaceGrid: - def __init__(self, side_pixels: int, unit_cell_pitch: float, prop_distances: np.ndarray, materials_refr_index: np.ndarray): + def __init__(self, side_pixels: int, unit_cell_pitch: float, prop_distances: np.ndarray, materials_refr_index: np.ndarray, 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 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.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) @@ -128,14 +134,20 @@ class SpaceGrid: def get_Y(self): return self.Y + def get_XY(self): + return self.X, 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 + def __init__(self, input_field: np.ndarray, wavelength: float, space_grid: SpaceGrid, center=(0,0)): + self.center = center self.space_grid = space_grid + self.input_field = self._shift_field(input_field) + self.wavelength = wavelength + + 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.") @@ -147,7 +159,16 @@ class 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"): + def get_center(self): + return self.center + + def _shift_field(self,input_field): + #take the input field and shift it according to the center + shift_field_x = np.roll(input_field, shift=int(self.center[0]/self.space_grid.get_sampling_period()), axis=1) + shift_field_y = np.roll(shift_field_x, shift=int(self.center[1]/self.space_grid.get_sampling_period()), axis=0) + return shift_field_y + + def plot2D(self, plot_zoom_x:float, plot_zoom_y:float, choose_quantity="amplitude", save_name="results/plot_out"): make_2Dplot_of(self.input_field, self.space_grid.X, self.space_grid.Y, @@ -157,17 +178,61 @@ class Field: 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) + def __init__(self, input_field: np.ndarray | str, wavelength: float, space_grid: SpaceGrid, wvl_unit_measure = 1, field_norm = False, center = (0,0)): + if isinstance(input_field, str): + input_field_to_constr = np.load(input_field) + input_field_shape = input_field_to_constr.shape + input_field_to_constr = (input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2))).reshape(input_field_shape) + else: + input_field_to_constr = input_field + + if field_norm: + input_field_to_constr = input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2)) + + super().__init__(input_field_to_constr, wavelength * wvl_unit_measure, space_grid, center=center) self.type = "source" def get_type(self): return self.type +class GaussianBeam(): + def __init__(self, beam_waist: float, space_grid: SpaceGrid, center=np.array([0,0]), field_norm = False): + self.beam_waist = beam_waist + self.center = center + self.space_grid = space_grid + generated_field = self._generate_gaussian_beam() + if field_norm: + self.field = generated_field / np.sqrt(np.sum(np.abs(generated_field)**2)) + else: + self.field = generated_field + + def _generate_gaussian_beam(self): + X = self.space_grid.X - self.center[0] + Y = self.space_grid.Y - self.center[1] + r_squared = X**2 + Y**2 + return np.exp(-r_squared / self.beam_waist**2) + + def get_field(self): + return self.field + + def get_space_grid(self): + return self.space_grid + 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) + def __init__(self, input_field: np.ndarray | str | Field ,associated_source: Source, wavelength: float, space_grid: SpaceGrid, wvl_unit_measure = 1, field_norm = False, center = (0,0)): + if isinstance(input_field, str): + input_field_to_constr = np.load(input_field) + input_field_shape = input_field_to_constr.shape + input_field_to_constr = (input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2))).reshape(input_field_shape) + else: + input_field_to_constr = input_field + + if field_norm: + input_field_to_constr = input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2)) + + super().__init__(input_field_to_constr, wavelength * wvl_unit_measure, space_grid, center=center) 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: @@ -232,7 +297,7 @@ class Simulation: 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()) @@ -284,12 +349,20 @@ class Simulation: 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() + mask_upper_bounds = np.concatenate((mask_upper_bounds, mask_upper_bounds),axis=0) + mask_lower_bounds = np.zeros_like(mask_upper_bounds) + + 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 = np.concatenate((mask_upper_bounds_symm, mask_upper_bounds_symm),axis=0) + mask_lower_bounds_symm = np.zeros_like(mask_upper_bounds_symm) + self.chosen_mask_upper_bounds = mask_upper_bounds_symm + self.chosen_mask_lower_bounds = mask_lower_bounds_symm + else: + self.chosen_mask_upper_bounds = mask_upper_bounds + self.chosen_mask_lower_bounds = mask_lower_bounds - 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 @@ -349,29 +422,59 @@ class Simulation: return w_symm.flatten() + def _build_circular_mask(self, center, radius): + ''' + Builds a circular mask centered at the given position with the specified radius. + ''' + x_center, y_center = center + X, Y = self.space_grid.get_XY() + rho = np.sqrt((X - x_center)**2 + (Y - y_center)**2) + return np.where(rho.flatten() < radius, 1, 0) + 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 ''' + source_to_initialize = self.sources[0] target_to_initialize = self.targets[0] + + source_Efield = source_to_initialize.get_field(flatten=True) + source_center = source_to_initialize.get_center() target_Efield = target_to_initialize.get_field(flatten=True) + target_center = target_to_initialize.get_center() 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 + circular_mask_target = self._build_circular_mask(target_center, self.phase_mask_init_size) + circular_mask_source = self._build_circular_mask(source_center, self.phase_mask_init_size) + source_pattern_phase = (np.angle(source_Efield) + 2 * np.pi) % (2 * np.pi) * circular_mask_source #make sure the phase is between 0 and 2pi + target_pattern_phase = (np.angle(target_Efield) + 2 * np.pi) % (2 * np.pi) * circular_mask_target #make sure the phase is between 0 and 2pi + + self.w_norm[:self.space_grid.get_S()] += source_pattern_phase / (2 * np.pi) #normalize the source phase to be between 0 and 1 self.w_norm[self.space_grid.get_S():] += target_pattern_phase / (2 * np.pi) #normalize the target phase to be between 0 and 1 + + 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():]) - 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.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) + #make a plot of the initialized phase mask + plt.imshow(self.phase_mask[:self.space_grid.get_S()].reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())), cmap='gray') + plt.title("Initialized Phase Mask") + plt.colorbar() + plt.savefig("results/initialized_phase_mask_1.pdf") + plt.imshow(self.phase_mask[self.space_grid.get_S():].reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())), cmap='gray') + plt.title("Initialized Phase Mask") + plt.colorbar() + plt.savefig("results/initialized_phase_mask_2.pdf") + + def _forward_propagate(self, source_object: Source): source_to_propagate = source_object.get_field() @@ -398,6 +501,10 @@ class Simulation: self.output_field = self.intermediate_fields[-1].flatten() + #make a plot of intermediate fields + #for i, field in enumerate(self.intermediate_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") + def _reduce_symmetrized_gradient(self, g_full): g_2d = g_full.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) @@ -439,6 +546,8 @@ class Simulation: wavelength = source.get_wavelength() self._update_wavelength(wavelength) + #we shall consider using a conic filter: before doing the forward propagation, + #we filter the masks self._forward_propagate(source) C_s = np.real(np.abs(np.sum(self.output_field * np.conj(target.get_field(flatten=True)))) ** 2) @@ -451,7 +560,6 @@ class Simulation: 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())) @@ -466,11 +574,14 @@ class Simulation: 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) + 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) + else: + ms1 = x[:self.space_grid.get_dof()//2] + ms2 = x[self.space_grid.get_dof()//2:] all_ms = np.concatenate((ms1, ms2)) @@ -487,9 +598,13 @@ class Simulation: 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) + if self.space_grid.xy_symmetry: + 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) + else: + reduced_grads = np.concatenate((full_grad1, full_grad2), axis=0) + grad_partial[:] += reduced_grads * 2 * np.pi / len(self.targets) #average the gradient across all wavelengths if grad.size > 0: @@ -498,7 +613,7 @@ class Simulation: #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: + if self.iter_num[0] % 5 == 0 and self.results_while_solving: fig, ax = plt.subplots() ax.plot(self.opt_history) ax.set_xlabel('Iteration') @@ -525,8 +640,20 @@ class Simulation: Placeholder method for exporting the results of the optimization. This method should be implemented to save the relevant data. ''' + #output fields + output_field = [] + for i,source in enumerate(self.sources): + self._forward_propagate(source) + output_field.append(self.intermediate_fields[-1]) + + np.save(f"{self.results_folder}/output_fields_{i}.npy", output_field[-1]) + make_2Dplot_of(output_field[-1], self.space_grid.get_X(), self.space_grid.get_Y(), plot_zoom_x=60, plot_zoom_y=60, choose_quantity="amplitude", save_name=f"{self.results_folder}/output_fields_{i}_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())) + make_2Dplot_of(phase_mask_1, self.space_grid.get_X(), self.space_grid.get_Y(), plot_zoom_x=60, plot_zoom_y=60, 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=60, plot_zoom_y=60, choose_quantity="phase_mask", save_name=f"{self.results_folder}/phase_mask_2_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) @@ -562,22 +689,26 @@ class Simulation: ''' 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_lower_bounds(self.chosen_mask_lower_bounds.flatten()) + solver.set_upper_bounds(self.chosen_mask_upper_bounds.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:]) + if self.space_grid.xy_symmetry: + 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:]) + else: + self.w_norm[:] = solver.optimize(self.w_norm) + self.is_optimized_flag = True + self.best_opt_val = solver.last_optimum_value() self._export_results() else: