diff --git a/filters.py b/filters.py new file mode 100644 index 0000000..4083713 --- /dev/null +++ b/filters.py @@ -0,0 +1,147 @@ +import autograd.numpy as np +from numpy.typing import ArrayLike as ArrayLikeType +import matplotlib.pyplot as plt + +filter_radius_um = 1 +def _centered(arr: np.ndarray, newshape: ArrayLikeType) -> np.ndarray: + """Formats the output of an FFT to center the zero-frequency component. + + A helper function borrowed from SciPy: + https://github.com/scipy/scipy/blob/v1.4.1/scipy/signal/signaltools.py#L263-L270 + + Args: + arr: output array from an FFT operation. + newshape: 1d array with two elements (integers) specifying the dimensions + of the array to be returned. + + Returns: + The input array with the zero-frequency component as the central element. + """ + newshape = np.asarray(newshape) + currshape = np.array(arr.shape) + startind = (currshape - newshape) // 2 + endind = startind + newshape + myslice = [slice(startind[k], endind[k]) for k in range(len(endind))] + + return arr[tuple(myslice)] + +def _edge_pad(arr: np.ndarray, pad: np.ndarray) -> np.ndarray: + """Border-pads the edges of an array. + + Used to preprocess the design weights prior to convolution with the filter. + Border padding an image will set the value of each padded pixel equal to + the value of the nearest pixel in the image. Used to implement feature- + preserving convolution filters that prevent unwanted edge effects. + + Args: + arr: 2d array whose borders contain the values to use for padding + pad: 2x2 array of integers indicating the size + of the borders to pad the array with. + + Returns: + A 2d array with border padding. + """ + # fill sides + left = np.tile(arr[0, :], (pad[0][0], 1)) + right = np.tile(arr[-1, :], (pad[0][1], 1)) + top = np.tile(arr[:, 0], (pad[1][0], 1)).transpose() + bottom = np.tile(arr[:, -1], (pad[1][1], 1)).transpose() + + # fill corners + top_left = np.tile(arr[0, 0], (pad[0][0], pad[1][0])) + top_right = np.tile(arr[-1, 0], (pad[0][1], pad[1][0])) + bottom_left = np.tile(arr[0, -1], (pad[0][0], pad[1][1])) + bottom_right = np.tile(arr[-1, -1], (pad[0][1], pad[1][1])) + + if pad[0][0] > 0 and pad[0][1] > 0 and pad[1][0] > 0 and pad[1][1] > 0: + return np.concatenate( + ( + np.concatenate((top_left, top, top_right)), + np.concatenate((left, arr, right)), + np.concatenate((bottom_left, bottom, bottom_right)), + ), + axis=1, + ) + elif pad[0][0] == 0 and pad[0][1] == 0 and pad[1][0] > 0 and pad[1][1] > 0: + return np.concatenate((top, arr, bottom), axis=1) + elif pad[0][0] > 0 and pad[0][1] > 0 and pad[1][0] == 0 and pad[1][1] == 0: + return np.concatenate((left, arr, right), axis=0) + elif pad[0][0] == 0 and pad[0][1] == 0 and pad[1][0] == 0 and pad[1][1] == 0: + return arr + else: + raise ValueError("At least one of the padding numbers is invalid.") + +def _quarter_to_full_kernel(arr: np.ndarray, pad_to: np.ndarray) -> np.ndarray: + """Constructs the full kernel from its nonnegative quadrant. + + Args: + arr: 2d input array representing the nonnegative quadrant of a + filter kernel with C4v symmetry. + pad_to: 1d array with two elements (integers) specifying the size + of the zero padding. + + Returns: + The complete kernel. + """ + pad_size = pad_to - 2 * np.array(arr.shape) + 1 + + top = np.zeros((pad_size[0], arr.shape[1])) + bottom = np.zeros((pad_size[0], arr.shape[1] - 1)) + middle = np.zeros((pad_to[0], pad_size[1])) + + top_left = arr[:, :] + top_right = np.flipud(arr[1:, :]) + bottom_left = np.fliplr(arr[:, 1:]) + bottom_right = np.flipud( + np.fliplr(arr[1:, 1:]) + ) # equivalent to flip, but flip is incompatible with autograd + + return np.concatenate( + ( + np.concatenate((top_left, top, top_right)), + middle, + np.concatenate((bottom_left, bottom, bottom_right)), + ), + axis=1, + ) + +def conic_filter(phase_mask,radius,space_grid): + '''Applies a conic filter to the input phase mask. The filter is defined as a + conic region of a specified radius, where the values are averaged. + + Args: + phase_mask: the input phase mask to be filtered (2D) + radius: the radius of the conic filter in micrometers + space_grid: an object containing the size and resolution of the phase mask + Returns: + The filtered phase mask (2D) + ''' + + Nx = space_grid.get_Nx() + Ny = space_grid.get_Ny() + + xv = np.arange(0, Nx//2-1) * (space_grid.get_resolution()) + yv = np.arange(0, Ny//2-1) * (space_grid.get_resolution()) + Xv, Yv = np.meshgrid(xv, yv) + + h = np.where(Xv**2 + Yv**2 < radius**2, (1 - np.sqrt(abs(Xv**2 + Yv**2)) / radius),0) + + h = _quarter_to_full_kernel(h, 3 * np.array([Nx, Ny])) + + h = h / np.sum(h) + #make a plot of the filter kernel + H = np.fft.fft2(h) + + phase_mask_2d = phase_mask.reshape(Nx,Ny) + phase_mask_padded = _edge_pad(phase_mask_2d, ((Nx, Nx), (Ny, Ny))) + + phase_mask_fft = np.fft.fft2(phase_mask_padded) + + Y = phase_mask_fft * H + + y = _centered( + np.real(np.fft.ifft2(Y)), + (Nx, Ny) + ) + + return y.flatten() \ No newline at end of file diff --git a/mpid.py b/mpid.py index 95f379c..af7970b 100644 --- a/mpid.py +++ b/mpid.py @@ -4,6 +4,8 @@ plt.rcParams.update({'font.size': 14}) import os import pyfftw import nlopt +from autograd import tensor_jacobian_product, grad as gd +from filters import conic_filter 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. @@ -110,6 +112,18 @@ class SpaceGrid: def get_sampling_period(self): return self.sampling_period + def get_resolution(self): + return self.unit_cell_pitch + + def get_unit_cell_pitch(self): + return self.unit_cell_pitch + + def get_size_x(self): + return self.size_x + + def get_size_y(self): + return self.size_y + def get_Nx(self): return self.Nx @@ -332,7 +346,7 @@ class Simulation: 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): + def set_optimization(self, opt_param_min, opt_param_max, ms_max_avail_diameter, phase_mask_init_size, opt_max_eval, conic_filter_radius, results_while_solving=True, custom_mask = None): ''' Sets the optimization parameters for the simulation. Args: @@ -348,6 +362,10 @@ class Simulation: 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 + self.filter_radius = conic_filter_radius + 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)) 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 = np.concatenate((mask_upper_bounds, mask_upper_bounds),axis=0) @@ -474,6 +492,14 @@ class Simulation: plt.colorbar() plt.savefig("results/initialized_phase_mask_2.pdf") + def _phase_filter(self, phase_mask): + '''Applies a conic filter on the phase mask''' + phase_mask_filtered = conic_filter( + phase_mask, + self.filter_radius, + self.space_grid, + ) + return phase_mask_filtered.flatten() def _forward_propagate(self, source_object: Source): @@ -486,14 +512,14 @@ class Simulation: 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_input_array[:,:] = np.fft.ifftshift((input_field * np.exp(1j * self.filtered_phase_mask[:self.space_grid.get_S()])).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))) self.fft_operator() self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[1] #element wise product self.ifft_operator() 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_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * np.exp(1j * self.filtered_phase_mask[self.space_grid.get_S():].reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())))) self.fft_operator() self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[2] #element wise product self.ifft_operator() @@ -520,8 +546,8 @@ 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.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())) + 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())) #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()) @@ -573,7 +599,6 @@ class Simulation: return w_full.flatten() def _cost_fun(self, x, grad): - 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) @@ -587,6 +612,12 @@ class Simulation: self.w_norm[:] = all_ms self.phase_mask[:] = self._phase_given_w(self.w_norm) + + + 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)) #after the phase mask is set, we test the same mask across all targets C_all_wavelengths = [] @@ -599,20 +630,25 @@ class Simulation: C_all_wavelengths.append(C) 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) + grad_partial_1 = self._reduce_symmetrized_gradient(full_grad1) + grad_partial_2 = self._reduce_symmetrized_gradient(full_grad2) else: - reduced_grads = np.concatenate((full_grad1, full_grad2), axis=0) + grad_partial_1 = full_grad1 + grad_partial_2 = full_grad2 - grad_partial[:] += reduced_grads * 2 * np.pi / len(self.targets) #average the gradient across all wavelengths + jac_grad_1 = tensor_jacobian_product(self._phase_filter)(self.phase_mask[:self.space_grid.get_S()], grad_partial_1) + jac_grad_2 = tensor_jacobian_product(self._phase_filter)(self.phase_mask[self.space_grid.get_S():], grad_partial_2) + jac_grad = np.concatenate((jac_grad_1, jac_grad_2),axis=0) + + grad_partial[:] += jac_grad if grad.size > 0: - grad[:] = grad_partial + grad[:] = grad_partial * 2 * np.pi / len(self.targets) #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 + print("Iteration:", self.iter_num[0]) if self.iter_num[0] % 5 == 0 and self.results_while_solving: fig, ax = plt.subplots() ax.plot(self.opt_history) diff --git a/testing_filtering.py b/testing_filtering.py new file mode 100644 index 0000000..ac4b205 --- /dev/null +++ b/testing_filtering.py @@ -0,0 +1,87 @@ +import mpid as pid +import numpy as np +import os +if not os.path.exists("results"): + os.makedirs("results") + +#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, +) + +wavelength = 1.55e-6 + +#define the sources +source1 = pid.Source(input_field='fiber_modes/LP11_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', + wavelength=wavelength, + space_grid=space_grid, + field_norm=True) + + +target = pid.Target(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy', + center=(20e-6,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=(-20e-6,0), + associated_source=source2, + 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], + central_wavelength=wavelength, + space_grid=space_grid, + n_cpu=1 +) + +opt_param_min = 0 +opt_param_max = 1 + +sim.set_optimization(opt_param_min=opt_param_min, + opt_param_max=opt_param_max, + ms_max_avail_diameter=ms_max_available_diameter, + phase_mask_init_size=output_fiber_core_diameter, + opt_max_eval=opt_max_eval, + conic_filter_radius=1.5e-6) + +sim.run() + +best_opt_val, distances = sim.get_best_opt_and_distances() +print("Best optimization value:", best_opt_val) +print("Distances:", distances) \ No newline at end of file