mirror of
https://github.com/edoverg/MPID.git
synced 2026-10-14 23:26:38 +00:00
first commit
This commit is contained in:
@@ -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.")
|
||||
|
||||
Reference in New Issue
Block a user