edoardo/MPID
1
0
Fork 0
mirror of https://github.com/edoverg/MPID.git synced 2026-10-14 23:26:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

several modification introduced: fixed dimensions to micrometers, introduced ego loop

Browse Source
This commit is contained in:
Edoardo Vergerio
2026-09-21 11:42:22 +02:00
parent 5f88780426
commit 31859271da
4 changed files with 776 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mpid.py
+199 -115
View File
@@ -56,9 +56,15 @@ def make_1Dplot_of(x_axis,y_axis,plot_zoom_x,save_name="plot_1D",) -> None:
raise ValueError("Lists x_axis and y_axis must have the same length.") raise ValueError("Lists x_axis and y_axis must have the same length.")
plt.figure() plt.figure()
for i in range(len(x_axis)): for i in range(len(x_axis)):
if x_axis[i] is None:
plt.plot(y_axis[i])
else:
plt.plot(x_axis[i], y_axis[i]) plt.plot(x_axis[i], y_axis[i])
plt.xlabel("x [um]") plt.xlabel("x [um]")
plt.tight_layout() plt.tight_layout()
if plot_zoom_x is None:
pass
else:
plt.xlim(-plot_zoom_x, plot_zoom_x) plt.xlim(-plot_zoom_x, plot_zoom_x)
plt.savefig(save_name + ".pdf") plt.savefig(save_name + ".pdf")
plt.close() plt.close()
@@ -68,7 +74,7 @@ class SpaceGrid:
Class representing the spatial grid for the simulation. Class representing the spatial grid for the simulation.
Spatial dimensions are intended to be measured in micrometers [um]. Spatial dimensions are intended to be measured in micrometers [um].
''' '''
def __init__(self, side_pixels: int, unit_cell_pitch: float, prop_distances: np.ndarray, materials_refr_index: np.ndarray, xy_symmetry=False): 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.Nx = side_pixels#for the moment we consider only square grids
self.Ny = side_pixels self.Ny = side_pixels
self.unit_cell_pitch = unit_cell_pitch * 1e-6 self.unit_cell_pitch = unit_cell_pitch * 1e-6
@@ -89,18 +95,14 @@ class SpaceGrid:
self.X, self.Y = np.meshgrid(self.xs, self.ys) self.X, self.Y = np.meshgrid(self.xs, self.ys)
self.rho = np.sqrt(self.X**2 + self.Y**2) self.rho = np.sqrt(self.X**2 + self.Y**2)
self.sampling_period = self.unit_cell_pitch self.sampling_period = self.unit_cell_pitch
if prop_distances is None:
self.propagation_distances_set = False
self.d = None
else:
self.propagation_distances_set = True
self.d = prop_distances * 1e-6 self.d = prop_distances * 1e-6
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 self.materials_refr_index = materials_refr_index
def get_grid_shape(self): def get_grid_shape(self):
@@ -161,18 +163,20 @@ class SpaceGrid:
return self.rho return self.rho
class Field: class Field:
def __init__(self, input_field: np.ndarray, wavelength: float, space_grid: SpaceGrid, center: np.ndarray = np.array([0,0])): def __init__(self, input_field: np.ndarray, wavelength: float, center: np.ndarray = np.array([0,0])):
self.center = center * 1e-6 self.center = center * 1e-6
self.space_grid = space_grid self.input_field = input_field
self.input_field = self._shift_field(input_field)
self.wavelength = wavelength * 1e-6 self.wavelength = wavelength * 1e-6
if self.input_field.shape != self.space_grid.get_grid_shape(): def get_field(self, space_grid: SpaceGrid = None, flatten=False):
raise ValueError("Input field must have the same shape as the space grid it is applied onto.") if space_grid is None:
raise ValueError("A SpaceGrid must be provided when invoking get_field()")
def get_field(self, flatten=False): self.space_grid = space_grid
if flatten: if flatten:
return self.input_field.flatten() return self._shift_field(self.input_field).flatten()
return self._shift_field(self.input_field)
def get_field_arr(self):
return self.input_field return self.input_field
def get_wavelength(self): def get_wavelength(self):
@@ -188,16 +192,17 @@ class Field:
return shift_field_y return shift_field_y
def plot2D(self, plot_zoom_x:float, plot_zoom_y:float, choose_quantity="amplitude", save_name="results/plot_out"): 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, return 0 #broken atm
self.space_grid.X, #make_2Dplot_of(self.input_field,
self.space_grid.Y, # self.space_grid.X,
plot_zoom_x, # self.space_grid.Y,
plot_zoom_y, # plot_zoom_x,
choose_quantity, # plot_zoom_y,
save_name) # choose_quantity,
# save_name)
class Source(Field): class Source(Field):
def __init__(self, input_field: np.ndarray | str, wavelength: float, space_grid: SpaceGrid, field_norm = False, center: np.ndarray = np.array([0,0])): def __init__(self, input_field: np.ndarray | str, wavelength: float, field_norm = False, center: np.ndarray = np.array([0,0])):
if isinstance(input_field, str): if isinstance(input_field, str):
input_field_to_constr = np.load(input_field) input_field_to_constr = np.load(input_field)
input_field_shape = input_field_to_constr.shape input_field_shape = input_field_to_constr.shape
@@ -208,7 +213,7 @@ class Source(Field):
if field_norm: if field_norm:
input_field_to_constr = input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2)) 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, space_grid, center=center) super().__init__(input_field_to_constr, wavelength, center=center)
self.type = "source" self.type = "source"
def get_type(self): def get_type(self):
@@ -241,7 +246,6 @@ class Target(Field):
def __init__(self, input_field: np.ndarray | str | Field , def __init__(self, input_field: np.ndarray | str | Field ,
associated_source: Source, associated_source: Source,
wavelength: float, wavelength: float,
space_grid: SpaceGrid,
field_norm = False, field_norm = False,
center = np.array([0,0])): center = np.array([0,0])):
if isinstance(input_field, str): if isinstance(input_field, str):
@@ -254,7 +258,7 @@ class Target(Field):
if field_norm: if field_norm:
input_field_to_constr = input_field_to_constr / np.sqrt(np.sum(np.abs(input_field_to_constr)**2)) 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, space_grid, center=center) super().__init__(input_field_to_constr, wavelength, center=center)
self.type = "target" self.type = "target"
if associated_source.get_wavelength() != self.get_wavelength(): if associated_source.get_wavelength() != self.get_wavelength():
@@ -300,48 +304,24 @@ class PlotArray():
raise ValueError("Input field must be 1D for plot1D.") raise ValueError("Input field must be 1D for plot1D.")
class Simulation: class Simulation:
def __init__(self, sources:list[Source], targets:list[Target], central_wavelength: float, space_grid: SpaceGrid, n_cpu: int = 1): def __init__(self, sources:list[Source],
targets:list[Target],
central_wavelength: float,
space_grid: SpaceGrid = None,
n_cpu: int = 1):
self.sources = sources self.sources = sources
self.targets = targets self.targets = targets
self.space_grid = space_grid self.eps_0 = 8.854e-12
self.central_wavelength = central_wavelength * 1e6 #central_wavelength parameter has little sense to exist. Revise self.mu_0 = 4*np.pi*1e-7
self.central_wavelength = central_wavelength * 1e-6 #central_wavelength parameter has little sense to exist. Revise
self.operating_wavelength = self.central_wavelength #initially, the operating_wavelength is the central wavelength 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.n_cpu = n_cpu
if space_grid is not None:
self.set_space_grid(space_grid)
self._configure_fftw(self.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_optimization_set = False
self.is_optimized_flag = False self.is_optimized_flag = False
@@ -391,30 +371,10 @@ class Simulation:
self.use_filtering = True self.use_filtering = True
self.filter_radius = None self.filter_radius = None
if self.use_filtering: if conic_filter_radius is not None:
self.filter_radius = conic_filter_radius * 1e-6 self.filter_radius = conic_filter_radius * 1e-6
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))
else: else:
self.filtered_phase_mask = self.phase_mask self.filter_radius = None
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_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_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
else:
self.chosen_mask_upper_bounds = mask_upper_bounds_full
self.chosen_mask_lower_bounds = mask_lower_bounds_full
self._solver_init()
def get_sources(self): def get_sources(self):
return self.sources return self.sources
@@ -452,9 +412,13 @@ class Simulation:
''' '''
assert new_wavelength > 0, "Wavelength must be positive." assert new_wavelength > 0, "Wavelength must be positive."
self.operating_wavelength = new_wavelength self.operating_wavelength = new_wavelength
self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()] self.lda_vector = [self.operating_wavelength / n for n in self.space_grid.get_materials_refr_index()]
self.eps_r = (self.space_grid.get_materials_refr_index())**2
self.k_vectors = [2 * np.pi / lda for lda in self.lda_vector] 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.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 = [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_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))] self.P_dagger_nat = [self.P_nat[i].T.conj() for i in range(len(self.P_nat))]
@@ -497,10 +461,10 @@ class Simulation:
is locally initialized around each source and target. is locally initialized around each source and target.
''' '''
source_Efields = [source.get_field(flatten=True) for source in self.sources] source_Efields = [source.get_field(space_grid=self.space_grid, flatten=True) for source in self.sources]
source_centers = [source.get_center() for source in self.sources] source_centers = [source.get_center() for source in self.sources]
target_Efields = [target.get_field(flatten=True) for target in self.targets] 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] 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((2*self.space_grid.get_S()))
@@ -546,6 +510,8 @@ class Simulation:
Returns: Returns:
ndarray: The filtered phase mask. ndarray: The filtered phase mask.
''' '''
if self.filter_radius is None:
raise ValueError("Filter radius is not set.")
phase_mask_cos = npa.cos(phase_mask) phase_mask_cos = npa.cos(phase_mask)
phase_mask_sin = npa.sin(phase_mask) phase_mask_sin = npa.sin(phase_mask)
@@ -565,35 +531,42 @@ class Simulation:
def _forward_propagate(self, source_object: Source): def _forward_propagate(self, source_object: Source):
#source field #source field
source_to_propagate = source_object.get_field() source_to_propagate = source_object.get_field(space_grid=self.space_grid)
self.fft_input_array[:,:] = np.fft.ifftshift(source_to_propagate.reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))) 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.fft_operator()
self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[0] self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[0]
self.ifft_operator() self.ifft_operator()
#field before MS1 #field before MS1
self.intermediate_fields[0] = np.fft.fftshift(self.ifft_operator.output_array) self.intermediate_electric_fields[0] = np.fft.fftshift(self.ifft_operator.output_array)
input_field = self.intermediate_fields[0].flatten()
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 * self.filtered_phase_mask[:self.space_grid.get_S()])).reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny())))
self.fft_operator() self.fft_operator()
self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[1] #element wise product self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[1] #element wise product
self.ifft_operator() self.ifft_operator()
#field before MS2 #field before MS2
self.intermediate_fields[1] = np.fft.fftshift(self.ifft_operator.output_array.copy()) self.intermediate_electric_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))
#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 * self.filtered_phase_mask[self.space_grid.get_S():].reshape((self.space_grid.get_Nx(), self.space_grid.get_Ny()))))
self.fft_operator() self.fft_operator()
self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[2] #element wise product self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_nat[2] #element wise product
self.ifft_operator() self.ifft_operator()
#output field
self.intermediate_fields[2] = np.fft.fftshift(self.ifft_operator.output_array.copy())
self.output_field = self.intermediate_fields[-1].flatten() #output field
self.intermediate_electric_fields[2] = np.fft.fftshift(self.ifft_operator.output_array.copy())
self.output_field = self.intermediate_electric_fields[-1].flatten()
#make a plot of intermediate fields #make a plot of intermediate fields
#for i, field in enumerate(self.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") # 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")
@@ -620,19 +593,19 @@ class Simulation:
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())) 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 #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_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.fft_operator()
self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[2] self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[2]
self.ifft_operator() self.ifft_operator()
self.fft_input_array[:,:] = np.fft.ifftshift(np.fft.fftshift(self.ifft_operator.output_array) * Phi_2_dagger) 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() 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.fft_operator()
self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[1] #element wise self.ifft_input_array[:,:] = self.fft_operator.output_array * self.P_dagger_nat[1] #element wise
self.ifft_operator() 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() 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()
return grad_C_phi_1, grad_C_phi_2 return grad_C_phi_1, grad_C_phi_2
@@ -644,11 +617,10 @@ class Simulation:
wavelength = source.get_wavelength() wavelength = source.get_wavelength()
self._update_wavelength(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) self._forward_propagate(source)
C_s = np.real(np.abs(np.sum(self.output_field * np.conj(target.get_field(flatten=True)))) ** 2) 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) grad_C_phi_1, grad_C_phi_2 = self._adjoint_propagate(target)
return C_s, grad_C_phi_1, grad_C_phi_2 return C_s, grad_C_phi_1, grad_C_phi_2
@@ -802,6 +774,7 @@ class Simulation:
#after the phase mask is set, we test the same mask across all targets #after the phase mask is set, we test the same mask across all targets
C_all_wavelengths = [] C_all_wavelengths = []
tested_wavelengths = []
grad_partial = np.zeros_like(grad) grad_partial = np.zeros_like(grad)
for target in self.targets: for target in self.targets:
@@ -809,7 +782,7 @@ class Simulation:
C, full_grad_1, full_grad_2 = self._value_and_grad(target) C, full_grad_1, full_grad_2 = self._value_and_grad(target)
C_all_wavelengths.append(C) C_all_wavelengths.append(C)
tested_wavelengths.append(target.get_wavelength())
if self.use_filtering: if self.use_filtering:
jac_grad_1 = tensor_jacobian_product(self._phase_filter)(phase_mask_1, full_grad_1) 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_2 = tensor_jacobian_product(self._phase_filter)(phase_mask_2, full_grad_2)
@@ -831,7 +804,19 @@ class Simulation:
#the new cost function considers the average overlap across all wavelengths. #the new cost function considers the average overlap across all wavelengths.
self.opt_history.append(np.mean(C_all_wavelengths)) self.opt_history.append(np.mean(C_all_wavelengths))
self._plot_history() self._plot_history()
if self.iter_num[0] % 5 == 0 and self.results_while_solving:
fig, ax = plt.subplots()
ax.plot(np.array(tested_wavelengths)*1e6, C_all_wavelengths, 'o-')
ax.set_xlabel('Wavelength [um]')
ax.set_ylabel('Overlap')
ax.set_title('Overlap Across Wavelengths')
fig.tight_layout()
fig.savefig(f"{self.results_folder}/overlap_across_wavelengths.pdf")
plt.close(fig)
return np.mean(C_all_wavelengths) return np.mean(C_all_wavelengths)
def _plot_history(self): def _plot_history(self):
@@ -854,7 +839,29 @@ class Simulation:
This method sets up the optimizer based on the chosen strategy. This method sets up the optimizer based on the chosen strategy.
''' '''
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))
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_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_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
else:
self.chosen_mask_upper_bounds = mask_upper_bounds_full
self.chosen_mask_lower_bounds = mask_lower_bounds_full
if self.strategy == 'mean': if self.strategy == 'mean':
self.solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof()) self.solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof())
self.solver.set_lower_bounds(self.chosen_mask_lower_bounds.flatten()) self.solver.set_lower_bounds(self.chosen_mask_lower_bounds.flatten())
@@ -864,7 +871,9 @@ class Simulation:
self.solver.set_maxeval(self.opt_max_eval) self.solver.set_maxeval(self.opt_max_eval)
elif self.strategy == 'epigraph': elif self.strategy == 'epigraph':
self.solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof() + 1) self.solver = nlopt.opt(nlopt.LD_CCSAQ, self.space_grid.get_dof() + 1)
epigraph_and_weights_lower_bounds = np.insert(self.chosen_mask_lower_bounds.flatten(), 0, -np.inf) epigraph_and_weights_lower_bounds = np.insert(self.chosen_mask_lower_bounds.flatten(), 0, -np.inf)
epigraph_and_weights_upper_bounds = np.insert(self.chosen_mask_upper_bounds.flatten(), 0, np.inf) epigraph_and_weights_upper_bounds = np.insert(self.chosen_mask_upper_bounds.flatten(), 0, np.inf)
@@ -894,11 +903,20 @@ class Simulation:
Placeholder method for exporting the results of the optimization. Placeholder method for exporting the results of the optimization.
This method should be implemented to save the relevant data. This method should be implemented to save the relevant data.
''' '''
#intermediate fields
#intermediate_field[0] is on MS1
#intermediate_field[1] is on MS2
#intermediate_field[2] is the output
for i,field in enumerate(self.intermediate_electric_fields):
np.save(f"{self.results_folder}/intermediate_electric_fields_{i}.npy", field)
make_2Dplot_of(field, 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="amplitude", save_name=f"{self.results_folder}/intermediate_electric_fields_{i}_plot")
make_2Dplot_of(field, 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}/intermediate_electric_fields_{i}_phase_plot")
#output fields #output fields
output_field = [] output_field = []
for i,source in enumerate(self.sources): for i,source in enumerate(self.sources):
self._forward_propagate(source) self._forward_propagate(source)
output_field.append(self.intermediate_fields[-1]) output_field.append(self.intermediate_electric_fields[-1])
np.save(f"{self.results_folder}/output_fields_{i}.npy", output_field[-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=1e6*self.get_ms_max_avail_diameter()/2, plot_zoom_y=1e6*self.get_ms_max_avail_diameter()/2, choose_quantity="amplitude", save_name=f"{self.results_folder}/output_fields_{i}_plot") 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="amplitude", save_name=f"{self.results_folder}/output_fields_{i}_plot")
@@ -938,22 +956,85 @@ class Simulation:
def is_optimized(self): def is_optimized(self):
return self.is_optimized_flag return self.is_optimized_flag
def update_space_grid(self, new_space_grid: SpaceGrid): def set_space_grid(self, new_space_grid: SpaceGrid):
'''
Update the space grid and associate parameters used for propagation.
'''
for source in self.sources:
if source.get_field_arr().shape != new_space_grid.get_grid_shape():
raise ValueError("The provided space grid is not compatible with the source fields.")
for target in self.targets:
if target.get_field_arr().shape != new_space_grid.get_grid_shape():
raise ValueError("The provided space grid is not compatible with the target fields.")
self.space_grid = new_space_grid
self.space_grid_is_set = False
self.is_optimized_flag = False self.is_optimized_flag = False
self.best_opt_val = 0 self.best_opt_val = 0
self.space_grid = new_space_grid
if self.space_grid.get_distances() is not None:
try:
if self.space_grid.get_materials_refr_index().shape != self.space_grid.get_distances().shape:
raise ValueError("Materials array must have the same shape as d array.")
except AttributeError:
raise ValueError("Materials must be an Numpy array.")
self.space_grid_is_set = True
self.d = self.space_grid.get_distances()
print("Using distances: ", self.d)
self.ks = 2 * np.pi / self.space_grid.get_sampling_period() 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.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.KX, self.KY = np.meshgrid(self.kappas, self.kappas)
self.K_parallel = np.sqrt(self.KX**2 + self.KY**2) self.K_parallel = np.sqrt(self.KX**2 + self.KY**2)
self.nu_parallel = self.K_parallel / (2 * np.pi) 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_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]
#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)
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',))
else:
self.space_grid_is_set = False
def update_propagation_distances(self, new_distances: np.ndarray):
'''
Update the propagation distances of the simulation by creating a new space grid
with the specified distances while maintaining the other properties of the current space grid.
'''
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,
prop_distances=new_distances,
materials_refr_index=self.space_grid.get_materials_refr_index(),
xy_symmetry=True,
)
self.set_space_grid(new_space_grid)
def run(self): def run(self):
''' """
Placeholder method for running the simulation. Run the optimization process for the simulation.
This method should be implemented with the actual simulation logic. This method checks if the space grid and optimization settings are set,
''' initializes the phase masks, and performs the optimization based on the chosen strategy.
"""
if self.space_grid_is_set:
if self.is_optimization_set: if self.is_optimization_set:
print("Initializing solver...")
self._solver_init()
print("Starting optimization...") print("Starting optimization...")
self._initialize_phase_masks() self._initialize_phase_masks()
#initialize nlopt solver #initialize nlopt solver
@@ -966,6 +1047,7 @@ class Simulation:
self.w_norm_symm[:] = self.solver.optimize(self.w_norm_symm) self.w_norm_symm[:] = self.solver.optimize(self.w_norm_symm)
else: else:
raise ValueError(f"Unknown strategy: {self.strategy}") raise ValueError(f"Unknown strategy: {self.strategy}")
self.is_optimized_flag = True self.is_optimized_flag = True
self.best_opt_val = self.solver.last_optimum_value() 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])
@@ -1009,7 +1091,9 @@ class Simulation:
if self.save_results: if self.save_results:
self._export_results() self._export_results()
self.iter_num[:] = [0]
return self.best_opt_val return self.best_opt_val
else: else:
raise RuntimeError("Optimization settings have not been specified.") raise RuntimeError("Optimization settings have not been specified.")
else:
raise RuntimeError("Space grid has not been specified.")
single_converter_symm.py
+24 -20
View File
@@ -5,20 +5,23 @@ if not os.path.exists("results"):
os.makedirs("results") os.makedirs("results")
#Space grid definition #Space grid definition
N = 1025 N = 1025
unit_cell_pitch = 700e-9 unit_cell_pitch = 0.700
n0 = 1.45 n0 = 1.45
n1 = 1.0 n1 = 1.0
n2 = 1.45 n2 = 1.45
mat_refr_index = np.array([n0, n1, n2]) mat_refr_index = np.array([n0, n1, n2])
d0 = 476e-6 #propagation distance: source - MS1 [m]
d1 = 510e-6 #propagation distance: MS1 - MS2 [m] #d0 = 400 #propagation distance: source - MS1 [um] NO FILT
d2 = d0 #propagation distance: MS2 - target [m] #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
d2 = d0 #propagation distance: MS2 - target [um]
d = np.array([d0, d1, d2]) d = np.array([d0, d1, d2])
ms_max_available_diameter = 122e-6 #[um] ms_max_available_diameter = 122 #[um]
output_fiber_core_diameter = 19e-6 #[um] output_fiber_core_diameter = 19 #[um]
opt_max_eval = 140 opt_max_eval = 140
space_grid = pid.SpaceGrid( space_grid = pid.SpaceGrid(
@@ -30,38 +33,39 @@ space_grid = pid.SpaceGrid(
) )
#define the source #define the source
beam_waist = 5.2e-6 beam_waist = 5.2
gaussian_field = np.exp(-(space_grid.get_rho())**2 / (beam_waist)**2) #Gaussian input field, flattened 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 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] wavelengths_to_test_um = [1.530,1.550,1.570] #[um]
central_wavelength_m = 1550 * space_grid.nm source_files = ["fiber_modes/LP01_field_distribution_lda_1530_N_1025.npy",
sources = [pid.Source(input_field=gaussian_field_normalized, "fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy",
"fiber_modes/LP01_field_distribution_lda_1570_N_1025.npy"]
sources = [pid.Source(input_field=np.load(file),
wavelength=wvl, wavelength=wvl,
space_grid=space_grid, field_norm=True) for file, wvl in zip(source_files, wavelengths_to_test_um)]
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') #sources[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/input_field')
#define targets #define targets
#import field data from files and create target objects #import field data from files and create target objects
target_files = ["fiber_modes/LP21_field_distribution_lda_1550_N_1025.npy",] 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_fields = [np.load(file) for file in target_files] target_fields = [np.load(file) for file in target_files]
targets = [pid.Target(input_field=target_fields[i], targets = [pid.Target(input_field=target_fields[i],
associated_source=sources[i], associated_source=sources[i],
wavelength=wvl, wavelength=wvl,
space_grid=space_grid, field_norm = True) for i, (file, wvl) in enumerate(zip(target_files, wavelengths_to_test_um))]
wvl_unit_measure=space_grid.nm,
field_norm = True) 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') #targets[0].plot2D(plot_zoom_x=38,plot_zoom_y=38,choose_quantity='amplitude', save_name='results/target_field')
#setup the simulation instance #setup the simulation instance
simulation = pid.Simulation( simulation = pid.Simulation(
sources=sources, sources=sources,
targets=targets, targets=targets,
central_wavelength=central_wavelength_m, central_wavelength=wavelengths_to_test_um[0],
space_grid=space_grid, space_grid=space_grid,
n_cpu=1 n_cpu=1
) )
@@ -72,7 +76,7 @@ simulation.set_optimization(
ms_max_avail_diameter=ms_max_available_diameter, ms_max_avail_diameter=ms_max_available_diameter,
phase_mask_init_size=output_fiber_core_diameter/3, phase_mask_init_size=output_fiber_core_diameter/3,
opt_max_eval=opt_max_eval, opt_max_eval=opt_max_eval,
conic_filter_radius=1.5e-6, conic_filter_radius=1.0,
) )
simulation.run() simulation.run()
single_converter_symm_ego.py
+479
View File
@@ -0,0 +1,479 @@
import mpid as pid
import numpy as np
from mpi4py.futures import MPIPoolExecutor
from smt.applications.ego import EGO, KRG
from smt.sampling_methods import LHS
from smt.design_space import (
DesignSpace,
)
from sklearn.model_selection import KFold, LeaveOneOut
from scipy.stats import norm
from scipy.optimize import minimize
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import matplotlib.animation as animation
import copy
import pickle
import os
import gc
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.45
n1 = 1.0
n2 = 1.45
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
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)
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"]
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,
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
def EI(GP, points, f_min):
pred = - GP.predict_values(points) #GP predicts overlaps (positive)
var = GP.predict_variances(points)
args0 = (f_min - pred) / np.sqrt(var)
args1 = (f_min - pred) * norm.cdf(args0)
args2 = np.sqrt(var) * norm.pdf(args0)
if var.size == 1 and var == 0.0: # can be use only if one point is computed
return 0.0
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)
loo = LeaveOneOut()
#available kernels: ‘matern52’, ‘matern32’ ‘abs_exp’, ‘squar_exp’
kernels_to_test = ['abs_exp', 'squar_exp', 'matern52', 'matern32']
all_z_values = []
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=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
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)
ax.grid(True, alpha=0.3)
ax.legend()
plt.savefig(f'{results_folder}/loo_standardized_residuals_gaussian_fit.pdf')
plt.close(fig)
best_kernel = kernels_to_test[np.argmin(np.array(all_std_values))]
#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()
#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)
#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
def run_ego(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
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
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)
for x_st in x_start
]
)
opt_success = opt_all[[opt_i["success"] for opt_i in opt_all]]
obj_success = np.array([opt_i["fun"] for opt_i in opt_success])
ind_min = np.argmin(obj_success) #in this case, the most negative -EI is the best (maximum EI)
#save the minimum EI value for this iteration to the history
EI_history.append(- obj_success[ind_min])
opt = opt_success[ind_min]
#select the non-optimal points for history purposes. Later, we will plot
#all the non-optimal points with crosses to show the search process of the n_start points
non_optimal_points = [opt_i["x"] for i, opt_i in enumerate(opt_all) if i != ind_min]
x_et_k = opt["x"] #this is the coordinate of the most interesting point: the one that maximizes the EI (minimize -EI)
d0_et_k = x_et_k[0]
d1_et_k = x_et_k[1]
d2_et_k = d0_et_k
d_et_k = np.round(np.array([d0_et_k, d1_et_k, d2_et_k]),0)
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
x_et_k_tested = np.atleast_2d(x_et_k)
x_et_k_tested_history.append(x_et_k_tested)
d0_EI = np.linspace(xlimits[0][0], xlimits[0][1], 100)
d1_EI = np.linspace(xlimits[1][0], xlimits[1][1], 100)
d0_EI_grid, d1_EI_grid = np.meshgrid(d0_EI, d1_EI)
EI_grid = np.zeros_like(d0_EI_grid)
for i in range(d0_EI_grid.shape[0]):
for j in range(d0_EI_grid.shape[1]):
EI_grid[i,j] = EI(sm, np.atleast_2d([d0_EI_grid[i,j], d1_EI_grid[i,j]]), f_min_k)
fig, ax = plt.subplots(figsize=(8,6))
contour = ax.contourf(d0_EI_grid, d1_EI_grid, EI_grid, levels=50, cmap='viridis')
fig.colorbar(contour, ax=ax, label='Expected Improvement')
#scatter the full history of tested points up to the current iteration k, the current point is red, the others are blue
x_et_k_tested_history_array = np.vstack(x_et_k_tested_history)
#scatter last point in red, the others in blue
ax.scatter(x_et_k_tested_history_array[:-1,0], x_et_k_tested_history_array[:-1,1], color='blue', label='Previous best')
ax.scatter(x_et_k_tested_history_array[-1,0], x_et_k_tested_history_array[-1,1], color='red', label='Best point')
#also plot the non optimal points with crosses for this iteration
if len(non_optimal_points) > 0:
non_optimal_points_array = np.vstack(non_optimal_points)
ax.scatter(non_optimal_points_array[:,0], non_optimal_points_array[:,1], color='orange', marker='x', label='Candidate')
ax.set_title(f'EGO Iteration {k+1}/{n_ego_iter}')
ax.set_xlabel('d0 [um]')
ax.set_ylabel('d1 [um]')
ax.set_xlim(xlimits[0])
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')
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.
ind_best = np.argmin(y_data)
x_opt = x_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")
im = plt.imshow(image_pt)
ims.append([im])
ani = animation.ArtistAnimation(fig, ims, interval=500)
#save the animation to a gif file
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')
ax.set_xlabel('EGO Iteration')
ax.set_ylabel('Expected Improvement')
ax.grid()
fig.tight_layout()
fig.savefig(f'{surrogate_folder_name}/EI_history.pdf')
plt.close(fig)
#surface plot of the refined surrogate model
d0_grid = np.linspace(xlimits[0][0], xlimits[0][1], 100)
d1_grid = np.linspace(xlimits[1][0], xlimits[1][1], 100)
d0_grid_mesh, d1_grid_mesh = np.meshgrid(d0_grid, d1_grid)
d_grid = np.column_stack((d0_grid_mesh.flatten(), d1_grid_mesh.flatten()))
sm_predictions = sm.predict_values(d_grid).reshape(d0_grid_mesh.shape)
fig = plt.figure(figsize=(8,6))
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_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]')
ax.legend()
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')
fig.colorbar(cp, ax=ax, label='Surrogate Model Prediction')
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.scatter([x_opt[0]], [x_opt[1]], color='green', s=100, label='Best Point', edgecolor='black')
ax.set_xlabel('d0 [um]')
ax.set_ylabel('d1 [um]')
ax.set_xlim(xlimits[0])
ax.set_ylim(xlimits[1])
ax.legend(framealpha=0.5)
fig.tight_layout()
fig.savefig(f'{surrogate_folder_name}/refined_surrogate_model_contour.pdf')
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'))
return x_opt, y_opt, ind_best, x_data, y_data
##############################
#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
#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'):
#build DOE
print("Building DOE...")
sampling = LHS(xlimits=xlimits, seed=seed)
xdoe = np.round(sampling(ndoe),0)
ydoe = []
ydoe = np.zeros((len(xdoe), 1), dtype=float)
for i, point in enumerate(xdoe):
print(f"DOE evaluation {i+1}/{len(xdoe)} for point {point}")
d0 = point[0]
d1 = point[1]
d2 = d0
d = np.round(np.array([d0, d1, d2]),0)
simulation.update_propagation_distances(d)
simulation.run()
ydoe[i, 0] = simulation.get_best_opt_val()
with open(f'{surrogate_folder_name}/ego_evaluations.txt', 'a') as f:
for xd, yd in zip(xdoe, ydoe):
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')
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)
trained_surrogate_model = leave_one_out(initial_dataset)
print("Starting EGO optimization...")
x_opt, y_opt, ind_best, x_data, y_data = run_ego(surrogate_model=trained_surrogate_model,
n_ego_iter=n_ego_iter,
xlimits=xlimits,
simulation=simulation)
sorter_to_gauss.py
+9 -9
View File
@@ -34,11 +34,11 @@ wavelength = 1.55
#define the sources #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/LP11_field_distribution_lda_1550_N_1025.npy',
wavelength=wavelength, wavelength=wavelength,
space_grid=space_grid, #space_grid=space_grid,
field_norm=True) 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/LP21_field_distribution_lda_1550_N_1025.npy',
wavelength=wavelength, wavelength=wavelength,
space_grid=space_grid, #space_grid=space_grid,
field_norm=True) field_norm=True)
@@ -46,20 +46,20 @@ target = pid.Target(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_
center=np.array([20,0]), center=np.array([20,0]),
associated_source=source1, associated_source=source1,
wavelength=wavelength, wavelength=wavelength,
space_grid=space_grid, #space_grid=space_grid,
field_norm=True) field_norm=True)
target2 = pid.Target(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy', target2 = pid.Target(input_field='fiber_modes/LP01_field_distribution_lda_1550_N_1025.npy',
center=np.array([-20,0]), center=np.array([-20,0]),
associated_source=source2, associated_source=source2,
wavelength=wavelength, wavelength=wavelength,
space_grid=space_grid, #space_grid=space_grid,
field_norm=True) field_norm=True)
#make a plot of sources and targets #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") #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") #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") #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") #source2.plot2D(plot_zoom_x=100, plot_zoom_y=100, choose_quantity="amplitude", save_name="results/source2_plot")
sim = pid.Simulation( sim = pid.Simulation(
sources=[source1, source2], sources=[source1, source2],
@@ -78,7 +78,7 @@ sim.set_optimization(opt_param_min=opt_param_min,
phase_mask_init_size=output_fiber_core_diameter, phase_mask_init_size=output_fiber_core_diameter,
opt_max_eval=opt_max_eval, opt_max_eval=opt_max_eval,
strategy = 'mean', strategy = 'mean',
conic_filter_radius = 1.0, #i.e. 700 [nm] #conic_filter_radius = 1.0, #i.e. 700 [nm]
) )
sim.run() sim.run()