mirror of
https://github.com/edoverg/MPID.git
synced 2026-10-14 23:26:38 +00:00
implemented conic filtering on phase masks. The purpose of filtering is to remove large spatial variations of phase within the metasurface. On the other hand, the disadvantage of filtering is that high spatial variations might be required to achieve the desired target. A compromise needs to be reached: choose a small enough radius to avoid spurious phase values but that can still yield good overlap values.
This commit is contained in:
Edoardo Vergerio
+147
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user