mirror of
https://github.com/edoverg/modeconverter.git
synced 2026-10-14 23:06:39 +00:00
432 lines
15 KiB
Python
432 lines
15 KiB
Python
#from jax import numpy as npa, grad
|
|
#from jax.numpy.fft import fft2, ifft2, fftshift, ifftshift
|
|
from typing import List, Tuple, Union
|
|
from autograd import numpy as npa, tensor_jacobian_product, grad as gd
|
|
import meep.adjoint as mpa
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import xarray as xr
|
|
from PIL import Image
|
|
import nlopt
|
|
ArrayLikeType = Union[List, Tuple, np.ndarray]
|
|
|
|
wavelength = 1.55
|
|
k0 = 2 * np.pi / wavelength
|
|
|
|
size_x = 300 * wavelength
|
|
size_y = 300 * wavelength
|
|
res_x = 3 #number of pixels per unit-length
|
|
res_y = res_x
|
|
|
|
Nx = int(res_x * size_x) + 1
|
|
Ny = int(res_y * size_y) + 1
|
|
S = Nx * Ny
|
|
|
|
norm_phase_min = 0
|
|
norm_phase_max = 1
|
|
w_norm = np.ones((S)) * 0.5 #normalized phase parameters in [0,1]
|
|
phase_min = 0
|
|
phase_max = 2 * np.pi
|
|
|
|
xs = np.linspace(-size_x/2, size_x/2, Nx)
|
|
ys = np.linspace(-size_y/2, size_y/2, Ny)
|
|
sampling_period = xs[1] - xs[0]
|
|
|
|
z_prop = 300 #propagation distance
|
|
|
|
input_field = np.ones((S))
|
|
|
|
opt_max_eval = 300
|
|
|
|
def phase_given_w(w):
|
|
'''Returns the phase shift given the design parameters w.
|
|
Args:
|
|
w: design parameters (flattened)
|
|
Returns:
|
|
The phase shift corresponding to the geometrical parameters w (flattened)'''
|
|
#for simplicity, we linearly map w (0,1) to a phase shift (0, 2pi)
|
|
return 2 * npa.pi * (w)
|
|
|
|
def get_pattern() -> np.ndarray:
|
|
'''Imports the intensity pattern from file'''
|
|
pattern_name = 'unipd.png'
|
|
pattern = Image.open('misc/' + pattern_name).convert('L')
|
|
pattern = np.array(pattern).astype(float)
|
|
|
|
pattern = pattern - np.min(pattern)
|
|
pattern = pattern / np.max(pattern)
|
|
|
|
return pattern
|
|
|
|
def intensity_desired_fn_pattern(
|
|
xs: list,
|
|
ys: list,
|
|
rescale: float = 0.5
|
|
) -> np.ndarray:
|
|
'''Returns the pattern intensity value at the specified (xs,ys) locations, with a scale factor.
|
|
Args:
|
|
xs: list of x coordinates
|
|
ys: list of y coordinates
|
|
rescale: scale factor for the pattern
|
|
Returns:
|
|
The flattened pattern intensity value at the specified (xs,ys) locations, with a scale factor
|
|
'''
|
|
pattern_values = get_pattern()
|
|
pattern_values = np.rot90(np.rot90(np.rot90(pattern_values)))
|
|
|
|
nx,ny = pattern_values.shape
|
|
xs_pattern = np.linspace(rescale*min(xs), rescale*max(xs), nx)
|
|
ys_pattern = np.linspace(rescale*min(ys), rescale*max(ys), ny)
|
|
pattern_dataArray = xr.DataArray(pattern_values, coords=dict(x=xs_pattern, y=ys_pattern))
|
|
pattern_interp = pattern_dataArray.interp(x=xs, y=ys)
|
|
|
|
pattern_final = npa.nan_to_num(pattern_interp.values, nan=npa.min(pattern_interp))
|
|
pattern_to_return = npa.nan_to_num(pattern_interp.values, nan=npa.min(pattern_interp)).flatten()
|
|
|
|
#make a plot before returning
|
|
plt.figure()
|
|
ax = plt.gca()
|
|
im = ax.imshow(
|
|
pattern_final,
|
|
extent=(min(xs), max(xs), min(ys), max(ys)),
|
|
origin='lower',
|
|
cmap='inferno',
|
|
)
|
|
ax.set_xlabel('x (um)')
|
|
ax.set_ylabel('y (um)')
|
|
ax.set_title('Target Intensity Pattern')
|
|
plt.colorbar(im, ax=ax, label='Intensity (a.u.)')
|
|
plt.savefig("results/test1_target_intensity_pattern.pdf")
|
|
plt.close()
|
|
return pattern_to_return
|
|
|
|
target_intensity = intensity_desired_fn_pattern(xs, ys)
|
|
mean_target_intensity = np.mean(target_intensity)
|
|
|
|
|
|
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 = npa.tile(arr[0, :], (pad[0][0], 1))
|
|
right = npa.tile(arr[-1, :], (pad[0][1], 1))
|
|
top = npa.tile(arr[:, 0], (pad[1][0], 1)).transpose()
|
|
bottom = npa.tile(arr[:, -1], (pad[1][1], 1)).transpose()
|
|
|
|
# fill corners
|
|
top_left = npa.tile(arr[0, 0], (pad[0][0], pad[1][0]))
|
|
top_right = npa.tile(arr[-1, 0], (pad[0][1], pad[1][0]))
|
|
bottom_left = npa.tile(arr[0, -1], (pad[0][0], pad[1][1]))
|
|
bottom_right = npa.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 npa.concatenate(
|
|
(
|
|
npa.concatenate((top_left, top, top_right)),
|
|
npa.concatenate((left, arr, right)),
|
|
npa.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 npa.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 npa.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 = npa.flipud(arr[1:, :])
|
|
bottom_left = npa.fliplr(arr[:, 1:])
|
|
bottom_right = npa.flipud(
|
|
npa.fliplr(arr[1:, 1:])
|
|
) # equivalent to flip, but flip is incompatible with autograd
|
|
|
|
return npa.concatenate(
|
|
(
|
|
npa.concatenate((top_left, top, top_right)),
|
|
middle,
|
|
npa.concatenate((bottom_left, bottom, bottom_right)),
|
|
),
|
|
axis=1,
|
|
)
|
|
|
|
def conic_filter(phase_mask,radius,size_x,size_y,resolution):
|
|
'''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
|
|
size_x: the size of the phase mask in x direction in micrometers
|
|
size_y: the size of the phase mask in y direction in micrometers
|
|
resolution: the resolution of the phase mask in pixels per micrometer
|
|
Returns:
|
|
The filtered phase mask (2D)
|
|
'''
|
|
Nx = int(round(size_x * resolution)) + 1
|
|
Ny = int(round(size_y * resolution)) + 1
|
|
|
|
xv = np.arange(0, size_x / 2, 1 / resolution) if resolution > 0 else [0]
|
|
yv = np.arange(0, size_y / 2, 1 / resolution) if resolution > 0 else [0]
|
|
|
|
X, Y = npa.meshgrid(xv, yv)
|
|
|
|
#h = npa.where(X**2 + Y**2 < radius**2, (1 - np.sqrt(abs(X**2 + Y**2)) / radius),0)
|
|
h = npa.where(X**2 + Y**2 < radius**2, 1, 0)
|
|
h = _quarter_to_full_kernel(h, 3 * np.array([Nx, Ny]))
|
|
h = h / npa.sum(h)
|
|
H = npa.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 = npa.fft.fft2(phase_mask_padded)
|
|
|
|
Y = phase_mask_fft * H
|
|
|
|
y = _centered(
|
|
npa.real(npa.fft.ifft2(Y)),
|
|
(Nx, Ny)
|
|
)
|
|
|
|
return y.flatten()
|
|
|
|
def phase_filter(phase_mask):
|
|
'''Applies a conic filter on the phase mask'''
|
|
phase_mask_filtered = conic_filter(
|
|
phase_mask.reshape(Nx,Ny), #row major mapping
|
|
filter_radius_um,
|
|
size_x,
|
|
size_y,
|
|
res_x,
|
|
)
|
|
return phase_mask_filtered.flatten()
|
|
|
|
def forward_propagate(phase_mask):
|
|
'''Propagates the input field, after applying the phase mask
|
|
|
|
Args:
|
|
phase_mask: the phase mask to be applied to the input field (flattened)
|
|
Returns:
|
|
A tuple with output_field (flattened) and corresponding propagation matrix P (2D)
|
|
'''
|
|
field_after_mask = input_field * npa.exp(1j * phase_mask)
|
|
field_after_mask_2d = field_after_mask.reshape((Nx, Ny))
|
|
field_fft = npa.fft.fftshift(npa.fft.fft2(field_after_mask_2d))
|
|
|
|
ks = 2 * npa.pi / sampling_period
|
|
kappas = npa.arange(-field_fft.shape[0]//2, field_fft.shape[0]//2) * ks / field_fft.shape[0]
|
|
KX, KY = npa.meshgrid(kappas, kappas)
|
|
K_parallel = npa.sqrt(KX**2 + KY**2)
|
|
nu_parallel = K_parallel / (2 * npa.pi)
|
|
phase_factor = k0 * z_prop * npa.sqrt(1 - (wavelength * nu_parallel) ** 2 + 0*1j)
|
|
|
|
P = npa.exp(1j * phase_factor)
|
|
fft_field_propagated = npa.multiply(field_fft,P) #element wise product
|
|
|
|
output_field_2d = npa.fft.ifft2(npa.fft.ifftshift(fft_field_propagated))
|
|
output_field = output_field_2d.flatten()
|
|
|
|
return output_field, P
|
|
|
|
def adjoint_propagate(x):
|
|
'''Backpropagates the output field using the adjoint of the propagation matrix P.
|
|
Args:
|
|
output_field: the output field to be backpropagated (flattened)
|
|
P: the propagation matrix used in the forward propagation
|
|
phase_mask: the phase mask used in the forward propagation (flattened)
|
|
Returns:
|
|
The backpropagated field (flattened)
|
|
'''
|
|
phase_mask = phase_given_w(x)
|
|
field_after_mask = input_field * npa.exp(1j * phase_mask)
|
|
field_after_mask_2d = field_after_mask.reshape((Nx, Ny))
|
|
field_fft = npa.fft.fftshift(npa.fft.fft2(field_after_mask_2d))
|
|
|
|
ks = 2 * npa.pi / sampling_period
|
|
kappas = npa.arange(-field_fft.shape[0]//2, field_fft.shape[0]//2) * ks / field_fft.shape[0]
|
|
KX, KY = npa.meshgrid(kappas, kappas)
|
|
K_parallel = npa.sqrt(KX**2 + KY**2)
|
|
nu_parallel = K_parallel / (2 * npa.pi)
|
|
phase_factor = k0 * z_prop * npa.sqrt(1 - (wavelength * nu_parallel) ** 2 + 0*1j)
|
|
|
|
P = npa.exp(1j * phase_factor)
|
|
fft_field_propagated = npa.multiply(field_fft,P) #element wise product
|
|
|
|
output_field_2d = npa.fft.ifft2(npa.fft.ifftshift(fft_field_propagated))
|
|
output_field = output_field_2d.flatten()
|
|
|
|
output_intensity = output_field * np.conj(output_field)
|
|
adjoint_field = (output_intensity - target_intensity) * output_field
|
|
adjoint_field_2d = adjoint_field.reshape((Nx, Ny))
|
|
adjoint_field_fft = npa.fft.fftshift(npa.fft.fft2(adjoint_field_2d))
|
|
P_dagger = P.T.conj()
|
|
adjoint_field_propagated_fft = adjoint_field_fft * P_dagger #element wise
|
|
adjoint_field_propagated_2d = npa.fft.ifft2(npa.fft.ifftshift(adjoint_field_propagated_fft))
|
|
|
|
#apply the conjugate of the phase mask to backpropagate through it
|
|
adjoint_field_propagated = adjoint_field_propagated_2d.flatten() * np.exp(-1j * phase_mask)
|
|
return adjoint_field_propagated
|
|
|
|
phase_mask_filtered = np.zeros((S))
|
|
def _compute_cost(phase_mask):
|
|
|
|
output_field, P = forward_propagate(phase_mask)
|
|
|
|
output_intensity = npa.real(output_field * npa.conj(output_field))
|
|
C_s = (output_intensity - target_intensity) ** 2
|
|
|
|
C = npa.sum(C_s)
|
|
|
|
return C
|
|
|
|
opt_history = []
|
|
def cost_fun(x, grad):
|
|
''' Cost function to be minimized.
|
|
Args:
|
|
x: input normalized parameter field (flattened)
|
|
grad: gradient of the cost function with respect to x, modified in place (flattened)
|
|
Returns:
|
|
The cost function value for the input parameter field x.
|
|
'''
|
|
w_norm[:] = x
|
|
|
|
phase_mask = phase_given_w(w_norm)
|
|
|
|
phase_mask_filtered = phase_filter(phase_mask)
|
|
|
|
cost_function_denominator = np.sum(np.ones((S)))
|
|
|
|
C = _compute_cost(phase_mask_filtered) / cost_function_denominator
|
|
opt_history.append(C)
|
|
#make a plot of the optimization history
|
|
plt.figure()
|
|
plt.plot(opt_history)
|
|
plt.xlabel('Iteration')
|
|
plt.ylabel('Cost Function Value')
|
|
plt.title('Optimization History')
|
|
plt.savefig("results/test1_optimization_history.pdf")
|
|
plt.close()
|
|
#first round of backpropagation
|
|
grad_test = gd(_compute_cost)(phase_mask_filtered)
|
|
|
|
#second round of backpropagation
|
|
full_grad = tensor_jacobian_product(phase_filter)(phase_mask,grad_test)
|
|
#print("grad from autograd",grad_test)
|
|
|
|
if grad.size > 0:
|
|
#adj_field_prop = adjoint_propagate(x)
|
|
#grad[:] = 4 * np.real(-1j * input_field.reshape((Nx,Ny)).T.conj().flatten() * adj_field_prop)
|
|
grad[:] = full_grad * 2 * np.pi / cost_function_denominator
|
|
#print("grad from theo",grad)
|
|
#print("grad diff",grad - grad_test)
|
|
return C
|
|
|
|
if __name__ == "__main__":
|
|
#initialize nlopt solver
|
|
solver = nlopt.opt(nlopt.LD_CCSAQ, S)
|
|
solver.set_lower_bounds(norm_phase_min)
|
|
solver.set_upper_bounds(norm_phase_max)
|
|
solver.set_min_objective(cost_fun)
|
|
solver.set_maxeval(opt_max_eval)
|
|
solver.set_param("dual_ftol_rel", 1e-7)
|
|
solver.set_param("verbosity",1)
|
|
|
|
print("Starting optimization...")
|
|
w_norm[:] = solver.optimize(w_norm)
|
|
print("Optimization completed.")
|
|
|
|
#verify the results
|
|
phase_mask = phase_given_w(w_norm)
|
|
filtered_mask = phase_filter(phase_mask)
|
|
#save the phase_mask as raw 2D text
|
|
np.savetxt("results/optimized_phase_mask.txt", filtered_mask.reshape((Nx, Ny)))
|
|
|
|
output_field, _ = forward_propagate(filtered_mask)
|
|
output_intensity = np.real(output_field * np.conj(output_field))
|
|
#make a plot of the phase mask, with phase wrap between -pi and pi
|
|
plt.figure()
|
|
plt.imshow(phase_mask.reshape((Nx, Ny)), extent=(-size_x/2, size_x/2, -size_y/2, size_y/2), origin='lower')
|
|
plt.colorbar()
|
|
plt.title('Phase Mask')
|
|
plt.xlabel('x (um)')
|
|
plt.ylabel('y (um)')
|
|
plt.savefig("results/test1_optimized_phase_mask.pdf")
|
|
plt.show()
|
|
plt.close()
|
|
#make a subplot to compare the output intensity with the target intensity
|
|
plt.figure(figsize=(12, 5))
|
|
plt.subplot(1, 2, 1)
|
|
plt.imshow(output_intensity.reshape((Nx, Ny)), extent=(-size_x/2, size_x/2, -size_y/2, size_y/2), origin='lower', cmap='inferno')
|
|
plt.colorbar()
|
|
plt.clim(0,1)
|
|
plt.title('Output Intensity')
|
|
plt.xlabel('x (um)')
|
|
plt.ylabel('y (um)')
|
|
plt.subplot(1, 2, 2)
|
|
plt.imshow(target_intensity.reshape((Nx, Ny)), extent=(-size_x/2, size_x/2, -size_y/2, size_y/2), origin='lower', cmap='inferno')
|
|
plt.colorbar()
|
|
plt.clim(0,1)
|
|
plt.title('Target Intensity')
|
|
plt.xlabel('x (um)')
|
|
plt.ylabel('y (um)')
|
|
plt.tight_layout()
|
|
plt.savefig("results/test1_comparison_output_target.pdf")
|
|
#plt.show()
|
|
plt.close() |