commit b0bfef5d27d316125f0a9d7d748868dc9aebda97 Author: Edoardo Vergerio Date: Thu Sep 3 16:59:35 2026 +0200 first commit diff --git a/mpid.py b/mpid.py new file mode 100644 index 0000000..e866612 --- /dev/null +++ b/mpid.py @@ -0,0 +1,585 @@ +import numpy as np +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 14}) +import os +import pyfftw +import nlopt + +def make_2Dplot_of(given_field,grid_X,grid_Y,plot_zoom_x,plot_zoom_y,choose_quantity="amplitude",save_name="plot") -> None: + '''Makes a plot of the given 2D field in a zoomed region. + Saves the plot with the chosen filename to a PDF file. + Args: + given_field: the 2D field to be plotted (flattened) + choose_quantity: either "amplitude" or "phase" to select quantity to plot + save_name: name of the file to save the plot (without extension) + plot_zoom_x: zoom level for the x-axis (in micrometers) + plot_zoom_y: zoom level for the y-axis (in micrometers) + Returns: + None''' + + given_field_2d = given_field.reshape((grid_X.shape[0], grid_X.shape[1])) #assure the field is 2D before plotting + fig, ax = plt.subplots(figsize=(6, 6)) + if choose_quantity == "amplitude": + ax.pcolormesh(grid_X*1e6,grid_Y*1e6,np.abs(given_field_2d),rasterized=True) + #add a colorbar for amplitude + plt.colorbar(ax.pcolormesh(grid_X*1e6,grid_Y*1e6,np.abs(given_field_2d),rasterized=True), ax=ax) + elif choose_quantity == "phase": + ax.pcolormesh(grid_X*1e6,grid_Y*1e6,(np.angle(given_field_2d)+2*np.pi)%(2*np.pi), cmap='hsv', vmin=0, vmax=2*np.pi, rasterized=True) + elif choose_quantity == "phase_mask": + ax.pcolormesh(grid_X*1e6,grid_Y*1e6,(given_field_2d+2*np.pi)%(2*np.pi), cmap='hsv', vmin=0, vmax=2*np.pi, rasterized=True) + else: + raise ValueError("Unexpected value for choose_quantity") + + ax.set_xlim(-plot_zoom_x, plot_zoom_x) + ax.set_ylim(-plot_zoom_y, plot_zoom_y) + ax.set_xlabel("x [um]") + ax.set_ylabel("y [um]") + plt.tight_layout() + plt.savefig(save_name + ".pdf") + plt.close() + +def make_1Dplot_of(x_axis,y_axis,plot_zoom_x,save_name="plot_1D",) -> None: + '''Makes a 1D plot of the given data and saves it to a PDF file. + Args: + x_axis: list of x-axis data (list of 1D arrays) + y_axis: list of y-axis data (list of 1D arrays) + plot_zoom_x: zoom level for the x-axis (in micrometers) + save_name: name of the file to save the plot (without extension) + Returns: + None + ''' + if len(x_axis) != len(y_axis): + raise ValueError("Lists x_axis and y_axis must have the same length.") + plt.figure() + for i in range(len(x_axis)): + plt.plot(x_axis[i], y_axis[i]) + plt.xlabel("x [um]") + plt.tight_layout() + plt.xlim(-plot_zoom_x, plot_zoom_x) + plt.savefig(save_name + ".pdf") + plt.close() + +class SpaceGrid: + def __init__(self, side_pixels: int, unit_cell_pitch: float, prop_distances: np.ndarray, materials_refr_index: np.ndarray): + self.Nx = side_pixels#for the moment we consider only square grids + self.Ny = side_pixels + self.unit_cell_pitch = unit_cell_pitch + self.nx = (self.Nx - 1) // 2 + self.ny = (self.Ny - 1) // 2 + self.dof = 2*((self.nx + 1)**2) #this needs to be computed depending on problem symmetry, which should be passed as argument (symmetry should be an instance of the Symmetry class) + self.S = self.Nx * self.Ny + self.size_x = self.unit_cell_pitch * self.Nx + self.size_y = self.unit_cell_pitch * self.Ny + self.xs = np.arange(-self.Nx//2+1, self.Nx//2+1) * (self.size_x / self.Nx) + self.ys = np.arange(-self.Ny//2+1, self.Ny//2+1) * (self.size_y / self.Ny) + self.X, self.Y = np.meshgrid(self.xs, self.ys) + self.rho = np.sqrt(self.X**2 + self.Y**2) + self.sampling_period = self.unit_cell_pitch + self.d = prop_distances + + self.nm = 1e-9 + self.m_to_nm = 1e9 + self.m = 1 + self.um = 1e-6 + self.m_to_um = 1e6 + try: + if materials_refr_index.shape != self.d.shape: + raise ValueError("Materials array must have the same shape as d array.") + except AttributeError: + raise ValueError("Materials must be an Numpy array.") + self.materials_refr_index = materials_refr_index + + def get_grid_shape(self): + return (self.Nx, self.Ny) + + def get_materials_refr_index(self): + return self.materials_refr_index + + def get_distances(self): + return self.d + + def get_num_distances(self): + return len(self.d) + + def get_sampling_period(self): + return self.sampling_period + + def get_Nx(self): + return self.Nx + + def get_Ny(self): + return self.Ny + + def get_nx(self): + return self.nx + + def get_ny(self): + return self.ny + + def get_dof(self): + return self.dof + + def get_S(self): + return self.S + + def get_X(self): + return self.X + + def get_Y(self): + return self.Y + + def get_rho(self): + return self.rho + +class Field: + def __init__(self, input_field: np.ndarray, wavelength: float, space_grid: SpaceGrid): + self.input_field = input_field + self.wavelength = wavelength + self.space_grid = space_grid + if self.input_field.shape != self.space_grid.get_grid_shape(): + raise ValueError("Input field must have the same shape as the space grid it is applied onto.") + + def get_field(self, flatten=False): + if flatten: + return self.input_field.flatten() + return self.input_field + + def get_wavelength(self): + return self.wavelength + + def plot2D(self, plot_zoom_x, plot_zoom_y, choose_quantity="amplitude", save_name="results/plot_out"): + make_2Dplot_of(self.input_field, + self.space_grid.X, + self.space_grid.Y, + plot_zoom_x, + plot_zoom_y, + choose_quantity, + save_name) + +class Source(Field): + def __init__(self, input_field: np.ndarray, wavelength: float, space_grid: SpaceGrid, wvl_unit_measure = 1,): + super().__init__(input_field, wavelength * wvl_unit_measure, space_grid) + self.type = "source" + + def get_type(self): + return self.type + +class Target(Field): + def __init__(self, input_field: np.ndarray, associated_source: Source, wavelength: float, space_grid: SpaceGrid, wvl_unit_measure = 1): + super().__init__(input_field, wavelength * wvl_unit_measure, space_grid) + self.type = "target" + if associated_source.get_wavelength() != wavelength * wvl_unit_measure: + raise ValueError("Associated source must have the same wavelength as the target.") + else: + self.associated_source = associated_source + + def get_type(self): + return self.type + + def get_associated_source(self): + return self.associated_source + +class PlotArray(): + def __init__(self, input_field: np.ndarray, space_grid: SpaceGrid): + self.input_field = input_field + self.array_shape = input_field.shape + self.space_grid = space_grid + self.type = "plot_array" + + def get_type(self): + return self.type + + def plot2D(self, plot_zoom_x, plot_zoom_y, choose_quantity="amplitude", save_name="results/plot2D_out"): + if len(self.array_shape) == 2: + make_2Dplot_of(self.input_field, + self.space_grid.X, + self.space_grid.Y, + plot_zoom_x, + plot_zoom_y, + choose_quantity, + save_name) + else: + raise ValueError("Input field must be 2D for plot2D.") + + def plot1D(self, plot_zoom_x, save_name="results/plot1D_out"): + if len(self.array_shape) == 1: + make_1Dplot_of([self.space_grid.X.flatten()], + [self.input_field.flatten()], + plot_zoom_x, + save_name) + else: + raise ValueError("Input field must be 1D for plot1D.") + +class Simulation: + def __init__(self, sources:list[Source], targets:list[Target], central_wavelength: float, space_grid: SpaceGrid, n_cpu: int = 1): + self.sources = sources + self.targets = targets + self.space_grid = space_grid + self.central_wavelength = central_wavelength #central_wavelength parameter has little sense to exist. Revise + + self.operating_wavelength = self.central_wavelength #initially, the operating_wavelength is the central wavelength + self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()] + + self.ks = 2 * np.pi / self.space_grid.get_sampling_period() + self.kappas = np.arange(-self.space_grid.get_Nx()//2+1, self.space_grid.get_Nx()//2+1) * self.ks / self.space_grid.get_Nx() + self.KX, self.KY = np.meshgrid(self.kappas, self.kappas) + self.K_parallel = np.sqrt(self.KX**2 + self.KY**2) + self.nu_parallel = self.K_parallel / (2 * np.pi) + + self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()] + self.k_vectors = [2 * np.pi / lda for lda in self.lda_vector] + self.phase_factors = [self.k_vectors[i] * self.space_grid.get_distances()[i] * np.sqrt(1 - (self.lda_vector[i] * self.nu_parallel) ** 2 + 0*1j) for i in range(len(self.k_vectors))] + + self.P = [np.exp(1j * self.phase_factors[i]) for i in range(len(self.phase_factors))] + self.P_nat = [np.fft.ifftshift(self.P[i]) for i in range(len(self.P))] + + #phase masks initialization + self.w_norm = np.zeros((2*self.space_grid.get_S())) + self.w_norm_symm = np.zeros(self.space_grid.get_dof()) + self.phase_mask = self._phase_given_w(self.w_norm) + #initialize intermediate fields + self.intermediate_fields = [None, None, None] + + #pyfftw init + self.n_cpu = n_cpu + self._configure_fftw(self.n_cpu) + self.fft_input_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16) + self.fft_output_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16) + self.ifft_input_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16) + self.ifft_output_array = pyfftw.empty_aligned((self.space_grid.get_Nx(), self.space_grid.get_Ny()), dtype='complex128', n=16) + self.fft_input_array.fill(0) + self.fft_output_array.fill(0) + self.ifft_input_array.fill(0) + self.ifft_output_array.fill(0) + self.fft_operator = pyfftw.FFTW(self.fft_input_array, self.fft_output_array, axes=(0,1), direction='FFTW_FORWARD', threads=self.n_cpu, flags=('FFTW_ESTIMATE',)) + self.ifft_operator = pyfftw.FFTW(self.ifft_input_array, self.ifft_output_array, axes=(0,1), direction='FFTW_BACKWARD', threads=self.n_cpu, flags=('FFTW_ESTIMATE',)) + + self.is_optimization_set = False + self.is_optimized_flag = False + self.results_while_solving = True + self.opt_history = [] + self.results_folder = "results" + if not os.path.exists(self.results_folder): + os.makedirs(self.results_folder) + + def _configure_fftw(self, n_cpu = 1, planner_effort='FFTW_ESTIMATE'): + pyfftw.config.NUM_THREADS = n_cpu + pyfftw.interfaces.cache.disable() + pyfftw.config.PLANNER_EFFORT = planner_effort + + def set_optimization(self, opt_param_min, opt_param_max, ms_max_avail_diameter, phase_mask_init_size, opt_max_eval, results_while_solving=True, custom_mask = None): + ''' + Sets the optimization parameters for the simulation. + Args: + optimization_params (dict): A dictionary containing optimization parameters. + ''' + self.is_optimization_set = True + #self.optimization_params = optimization_params + self.iter_num = [0] + self.norm_phase_min = opt_param_min + self.norm_phase_max = opt_param_max + self.ms_max_avail_diameter = ms_max_avail_diameter + self.phase_mask_init_size = phase_mask_init_size + self.opt_max_eval = opt_max_eval + self.results_while_solving = results_while_solving + self.custom_mask = custom_mask #future: one might be able to directly pass a custom mask + + + mask_upper_bounds = self.norm_phase_max * np.where( self.space_grid.get_X().flatten()**2 + self.space_grid.get_Y().flatten()**2 < (self.ms_max_avail_diameter/2)**2, 1, 0) + mask_upper_bounds_symm = mask_upper_bounds.reshape((self.space_grid.get_Nx(),self.space_grid.get_Ny()))[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1].flatten() + + self.mask_upper_bounds_symm = np.concatenate((mask_upper_bounds_symm, mask_upper_bounds_symm),axis=0) + self.mask_lower_bounds_symm = np.zeros_like(self.mask_upper_bounds_symm) + + def get_sources(self): + return self.sources + + def get_targets(self): + return self.targets + + def get_space_grid(self): + return self.space_grid + + def get_central_wavelength(self): + return self.central_wavelength + + def get_operating_wavelength(self): + return self.operating_wavelength + + def get_k_vectors(self): + return self.k_vectors + + def get_wvl_all_targets(self, unit_measure = 1): + return [target.get_wavelength() * unit_measure for target in self.targets] + + def _update_wavelength(self, new_wavelength): + ''' + Updates the operating wavelength and related vectors (lda_vector and k_vectors) based on the new wavelength. + The new_wavelength should be provided in meters. + Args: + new_wavelength (float): The new operating wavelength in meters. + ''' + assert new_wavelength > 0, "Wavelength must be positive." + self.operating_wavelength = new_wavelength + self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()] + self.k_vectors = [2 * np.pi / lda for lda in self.lda_vector] + self.phase_factors = [self.k_vectors[i] * self.space_grid.get_distances()[i] * np.sqrt(1 - (self.lda_vector[i] * self.nu_parallel) ** 2 + 0*1j) for i in range(len(self.k_vectors))] + self.P = [np.exp(1j * self.phase_factors[i]) for i in range(len(self.phase_factors))] + self.P_nat = [np.fft.ifftshift(self.P[i]) for i in range(len(self.P))] + self.P_dagger_nat = [self.P_nat[i].T.conj() for i in range(len(self.P_nat))] + + def _phase_given_w(self): + return 2 * np.pi * self.w_norm + + def _phase_given_w(self, input_normalized_phase): + return 2 * np.pi * input_normalized_phase + + def _symmetrize(self, w): + ''' + Takes an input array, which must be full in size, turns it into a 2D array, and symmetrizes it along the x and y axes. + Returns the first quadrant of shape (nx,ny), including the zero axis, flattened + ''' + try: + w_2d = w.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + except: + raise ValueError(f"Wrong shape for input array. Current shape: {w.shape}") + + w_symm = np.zeros((self.space_grid.get_nx()+1, self.space_grid.get_ny()+1)) + w_symm[:,:] = w_2d[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1] + + return w_symm.flatten() + + def _initialize_phase_masks(self): + ''' + Initializes the phase masks for the simulation. + This method should be called once, before iterating over the targets. + If multiple targets are present in this simulation instance, the first initialization uses the first target + ''' + target_to_initialize = self.targets[0] + target_Efield = target_to_initialize.get_field(flatten=True) + + self.w_norm = np.zeros((2*self.space_grid.get_S())) + + circular_mask_1 = np.where(self.space_grid.get_rho().flatten() < self.phase_mask_init_size, 1, 0) + target_pattern_phase = (np.angle(target_Efield) + 2 * np.pi) % (2 * np.pi)*circular_mask_1 #make sure the phase is between 0 and 2pi + + self.w_norm[self.space_grid.get_S():] += target_pattern_phase / (2 * np.pi) #normalize the target phase to be between 0 and 1 + + w_norm_symm_1 = self._symmetrize(self.w_norm[:self.space_grid.get_S()]) + w_norm_symm_2 = self._symmetrize(self.w_norm[self.space_grid.get_S():]) + + self.w_norm_symm[:] = np.concatenate((w_norm_symm_1, w_norm_symm_2), axis=0) + + self.phase_mask[:] = self._phase_given_w(self.w_norm) + + def _forward_propagate(self, source_object: Source): + + source_to_propagate = source_object.get_field() + + self.fft_input_array[:,:] = np.fft.ifftshift(source_to_propagate.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))) + self.fft_operator() + self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[0] + self.ifft_operator() + self.intermediate_fields[0] = np.fft.fftshift(self.ifft_operator.output_array) + input_field = self.intermediate_fields[0].flatten() + + self.fft_input_array[:,:] = np.fft.ifftshift((input_field * np.exp(1j * self.phase_mask[:self.space_grid.get_S()])).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))) + self.fft_operator() + self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[1] #element wise product + self.ifft_operator() + self.intermediate_fields[1] = np.fft.fftshift(self.ifft_operator.output_array.copy()) + #print("Field intensity after first phase mask and propagation: %.4e" % np.sum(np.abs(self.intermediate_fields[0])**2)) + + self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * np.exp(1j * self.phase_mask[self.space_grid.get_S():].reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())))) + self.fft_operator() + self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[2] #element wise product + self.ifft_operator() + self.intermediate_fields[2] = np.fft.fftshift(self.ifft_operator.output_array.copy()) + + self.output_field = self.intermediate_fields[-1].flatten() + + + def _reduce_symmetrized_gradient(self, g_full): + g_2d = g_full.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + + third = g_2d[self.space_grid.get_nx()+1:,:self.space_grid.get_ny()+1] + second = g_2d[:self.space_grid.get_nx()+1, self.space_grid.get_ny()+1:] + fourth = g_2d[self.space_grid.get_nx()+1:, self.space_grid.get_ny()+1:] + + g_2d[:self.space_grid.get_nx(), self.space_grid.get_ny()+1:] += np.flip(fourth, axis = 0) + g_2d[:self.space_grid.get_nx(), :self.space_grid.get_ny()+1] += np.flip(third, axis = 0) + g_2d[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()] += np.flip(second, axis = 1) + + return g_2d[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1].flatten() + + def _adjoint_propagate(self, target): + Phi_1_dagger = np.exp(-1j * self.phase_mask[:self.space_grid.get_S()]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + Phi_2_dagger = np.exp(-1j * self.phase_mask[self.space_grid.get_S():]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + + #compute the adjoint source + self.fft_input_array[:,:] = np.fft.ifftshift(np.sum(self.intermediate_fields[2] * np.conj(target.get_field())) * target.get_field()) + self.fft_operator() + self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[2] + self.ifft_operator() + self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * Phi_2_dagger) + + grad_C_phi_2 = 2 * np.real(-1j * self.intermediate_fields[1].conj() * np.fft.fftshift(self.fft_input_array)).flatten() + + self.fft_operator() + self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[1] #element wise + self.ifft_operator() + + grad_C_phi_1 = 2 * np.real(-1j * self.intermediate_fields[0].conj() * np.fft.fftshift(self.ifft_operator.output_array) * Phi_1_dagger).flatten() + + return grad_C_phi_1, grad_C_phi_2 + + def _value_and_grad(self, target): + + source = target.get_associated_source() + wavelength = source.get_wavelength() + + self._update_wavelength(wavelength) + self._forward_propagate(source) + + C_s = np.real(np.abs(np.sum(self.output_field * np.conj(target.get_field(flatten=True)))) ** 2) + grad_C_phi_1, grad_C_phi_2 = self._adjoint_propagate(target) + + return C_s, grad_C_phi_1, grad_C_phi_2 + + def _expand_symmetrize(self, w_symm): + ''' + Takes a array of shape (nx+1, ny+1), representing the first quadrant, including zero axes, and + expands it to a full 2D array of shape (Nx, Ny) by mirroring, flattened. + ''' + + w_symm_2d = w_symm.reshape((self.space_grid.get_nx()+1, self.space_grid.get_ny()+1)) + + w_full = np.zeros((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + + w_full[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()+1] = w_symm_2d + w_full[:self.space_grid.get_nx()+1, self.space_grid.get_ny()+1:] = np.flip(w_full[:self.space_grid.get_nx()+1, :self.space_grid.get_ny()], axis=1) + w_full[self.space_grid.get_nx()+1:, :self.space_grid.get_ny()+1] = np.flip(w_full[:self.space_grid.get_nx(), :self.space_grid.get_ny()+1], axis=0) + w_full[self.space_grid.get_nx()+1:, self.space_grid.get_ny()+1:] = np.flip(np.flip(w_full[:self.space_grid.get_nx(), :self.space_grid.get_ny()], axis=1), axis=0) + w_full[self.space_grid.get_nx(), self.space_grid.get_ny()] = w_symm_2d[-1, -1] + + return w_full.flatten() + + def _cost_fun(self, x, grad): + + ms1_first_quadrant = x[:self.space_grid.get_dof()//2].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + ms2_first_quadrant = x[self.space_grid.get_dof()//2:].reshape(self.space_grid.get_Nx()//2+1, self.space_grid.get_Ny()//2+1) + + ms1 = self._expand_symmetrize(ms1_first_quadrant) + ms2 = self._expand_symmetrize(ms2_first_quadrant) + + all_ms = np.concatenate((ms1, ms2)) + + self.w_norm[:] = all_ms + self.phase_mask[:] = self._phase_given_w(self.w_norm) + + #after the phase mask is set, we test the same mask across all targets + C_all_wavelengths = [] + grad_partial = np.zeros_like(grad) + + for target in self.targets: + + C, full_grad1, full_grad2 = self._value_and_grad(target) + + C_all_wavelengths.append(C) + + reduced_grad1 = self._reduce_symmetrized_gradient(full_grad1) + reduced_grad2 = self._reduce_symmetrized_gradient(full_grad2) + reduced_grads = np.concatenate((reduced_grad1, reduced_grad2), axis=0) + grad_partial[:] += reduced_grads * 2 * np.pi / len(self.targets) #average the gradient across all wavelengths + + if grad.size > 0: + grad[:] = grad_partial + + #the new cost function considers the average overlap across all wavelengths. + self.opt_history.append(np.mean(C_all_wavelengths)) + self.iter_num[0] += 1 + if self.iter_num[0] % 10 == 0 and self.results_while_solving: + fig, ax = plt.subplots() + ax.plot(self.opt_history) + ax.set_xlabel('Iteration') + ax.set_ylabel('Cost Function Value') + ax.set_title('Optimization History') + fig.tight_layout() + fig.savefig(f"{self.results_folder}/optimization_history.pdf") + plt.close(fig) + + #make a plot of the cost function values across + # all wavelengths for this iteration + fig, ax = plt.subplots() + ax.plot(self.get_wvl_all_targets(unit_measure=self.space_grid.m_to_nm), C_all_wavelengths, marker='o') + ax.set_xlabel(f'Wavelength [nm]') + ax.set_ylabel('Cost function value') + fig.tight_layout() + fig.savefig(f"{self.results_folder}/cost_function_across_wavelengths.pdf") + plt.close(fig) + + return np.mean(C_all_wavelengths) + + def _export_results(self): + ''' + Placeholder method for exporting the results of the optimization. + This method should be implemented to save the relevant data. + ''' + phase_mask_1 = self._phase_given_w(self.w_norm[:self.space_grid.get_S()]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + phase_mask_2 = self._phase_given_w(self.w_norm[self.space_grid.get_S():]).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())) + np.save(f"{self.results_folder}/optimized_phase_mask_1.npy", phase_mask_1) + np.save(f"{self.results_folder}/optimized_phase_mask_2.npy", phase_mask_2) + + def get_best_opt_val(self): + if self.is_optimized_flag: + return self.best_opt_val + else: + raise RuntimeError("Simulation object has not been optimized") + + def get_best_opt_and_distances(self): + if self.is_optimized_flag: + return self.best_opt_val, self.space_grid.d + else: + raise RuntimeError("Simulation object has not been optimized") + + def is_optimized(self): + return self.is_optimized_flag + + def update_space_grid(self, new_space_grid: SpaceGrid): + self.is_optimized_flag = False + self.best_opt_val = 0 + self.space_grid = new_space_grid + self.ks = 2 * np.pi / self.space_grid.get_sampling_period() + self.kappas = np.arange(-self.space_grid.get_Nx()//2+1, self.space_grid.get_Nx()//2+1) * self.ks / self.space_grid.get_Nx() + self.KX, self.KY = np.meshgrid(self.kappas, self.kappas) + self.K_parallel = np.sqrt(self.KX**2 + self.KY**2) + self.nu_parallel = self.K_parallel / (2 * np.pi) + + def run(self): + ''' + Placeholder method for running the simulation. + This method should be implemented with the actual simulation logic. + ''' + if self.is_optimization_set: + print("Starting optimization...") + + self._initialize_phase_masks() + #initialize nlopt solver + solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof()) + solver.set_lower_bounds(self.mask_lower_bounds_symm.flatten()) + solver.set_upper_bounds(self.mask_upper_bounds_symm.flatten()) + solver.set_max_objective(self._cost_fun) + solver.set_maxeval(self.opt_max_eval) + solver.set_param("dual_ftol_rel", 1e-7) + solver.set_param("verbosity",0) + + self.w_norm_symm[:] = solver.optimize(self.w_norm_symm) + self.is_optimized_flag = True + self.best_opt_val = solver.last_optimum_value() + self.w_norm[:self.space_grid.get_S()] = self._expand_symmetrize(self.w_norm_symm[:self.space_grid.get_dof()//2]) + self.w_norm[self.space_grid.get_S():] = self._expand_symmetrize(self.w_norm_symm[self.space_grid.get_dof()//2:]) + + self._export_results() + else: + raise RuntimeError("Optimization settings have not been specified.") + \ No newline at end of file diff --git a/testing.py b/testing.py new file mode 100644 index 0000000..750c9f6 --- /dev/null +++ b/testing.py @@ -0,0 +1,78 @@ +import mpid as pid +import numpy as np + +#Space grid definition +N = 1025 +unit_cell_pitch = 700e-9 + +n0 = 1.45 +n1 = 1.0 +n2 = 1.45 +mat_refr_index = np.array([n0, n1, n2]) + +d0 = 476e-6 #propagation distance: source - MS1 [m] +d1 = 510e-6 #propagation distance: MS1 - MS2 [m] +d2 = d0 #propagation distance: MS2 - target [m] +d = np.array([d0, d1, d2]) + +ms_max_available_diameter = 122e-6 #[um] +output_fiber_core_diameter = 19e-6 #[um] +opt_max_eval = 140 + +space_grid = pid.SpaceGrid( + side_pixels=N, + unit_cell_pitch = unit_cell_pitch, + prop_distances=d, + materials_refr_index=mat_refr_index, +) + +#define the source +beam_waist = 5.2e-6 +gaussian_field = np.exp(-(space_grid.get_rho())**2 / (beam_waist)**2) #Gaussian input field, flattened +gaussian_field_normalized = (gaussian_field / np.sqrt(np.sum(np.abs(gaussian_field)**2))) #normalize the input field to have power = 1 + +wavelengths_to_test_nm = [1550] #[nm] +central_wavelength_m = 1550 * space_grid.nm +sources = [pid.Source(input_field=gaussian_field_normalized, + wavelength=wvl, + space_grid=space_grid, + wvl_unit_measure=space_grid.nm) for wvl in wavelengths_to_test_nm] + +sources[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/input_field') +#define targets +#import field data from files and create target objects +target_files = ["fiber_modes/LP21_field_distribution_1025.npy"] + +target_fields = [np.load(file) for file in target_files] +norm_target_fields = [field / np.sqrt(np.sum(np.abs(field)**2)) for field in target_fields] + +targets = [pid.Target(input_field=norm_target_fields[i], + associated_source=sources[i], + wavelength=wvl, + space_grid=space_grid, + wvl_unit_measure=space_grid.nm) for i, (file, wvl) in enumerate(zip(target_files, wavelengths_to_test_nm))] + +targets[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/target_field') + +#setup the simulation instance +simulation = pid.Simulation( + sources=sources, + targets=targets, + central_wavelength=central_wavelength_m, + space_grid=space_grid, + n_cpu=1 +) + +simulation.set_optimization( + opt_param_min=0, + opt_param_max=1, + ms_max_avail_diameter=ms_max_available_diameter, + phase_mask_init_size=output_fiber_core_diameter/3, + opt_max_eval=opt_max_eval, +) + +simulation.run() + +best_opt_val, distances = simulation.get_best_opt_and_distances() +print("Best optimization value:", best_opt_val) +print("Distances:", distances) \ No newline at end of file