implemented and did some testing about applying a conic filter on phase masks. This should smooth out outliers at the cost of a possible reduction of total overlap efficiency. Several example devices have been added for testing purposes.

This commit is contained in:
Edoardo Vergerio
2026-09-08 16:56:56 +02:00
parent b4e79330b1
commit 45492233d4
5 changed files with 250 additions and 108 deletions
+44 -42
View File
@@ -1,9 +1,9 @@
import autograd.numpy as np
import autograd.numpy as npa
from numpy.typing import ArrayLike as ArrayLikeType
import matplotlib.pyplot as plt
filter_radius_um = 1
def _centered(arr: np.ndarray, newshape: ArrayLikeType) -> np.ndarray:
def _centered(arr: npa.ndarray, newshape: ArrayLikeType) -> npa.ndarray:
"""Formats the output of an FFT to center the zero-frequency component.
A helper function borrowed from SciPy:
@@ -17,15 +17,15 @@ def _centered(arr: np.ndarray, newshape: ArrayLikeType) -> np.ndarray:
Returns:
The input array with the zero-frequency component as the central element.
"""
newshape = np.asarray(newshape)
currshape = np.array(arr.shape)
newshape = npa.asarray(newshape)
currshape = npa.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:
def _edge_pad(arr: npa.ndarray, pad: npa.ndarray) -> npa.ndarray:
"""Border-pads the edges of an array.
Used to preprocess the design weights prior to convolution with the filter.
@@ -42,36 +42,36 @@ def _edge_pad(arr: np.ndarray, pad: np.ndarray) -> np.ndarray:
A 2d array with border padding.
"""
# fill sides
left = np.tile(arr[0, :], (pad[0][0], 1))
right = np.tile(arr[-1, :], (pad[0][1], 1))
top = np.tile(arr[:, 0], (pad[1][0], 1)).transpose()
bottom = np.tile(arr[:, -1], (pad[1][1], 1)).transpose()
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 = np.tile(arr[0, 0], (pad[0][0], pad[1][0]))
top_right = np.tile(arr[-1, 0], (pad[0][1], pad[1][0]))
bottom_left = np.tile(arr[0, -1], (pad[0][0], pad[1][1]))
bottom_right = np.tile(arr[-1, -1], (pad[0][1], pad[1][1]))
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 np.concatenate(
return npa.concatenate(
(
np.concatenate((top_left, top, top_right)),
np.concatenate((left, arr, right)),
np.concatenate((bottom_left, bottom, bottom_right)),
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 np.concatenate((top, arr, bottom), axis=1)
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 np.concatenate((left, arr, right), axis=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:
def _quarter_to_full_kernel(arr: npa.ndarray, pad_to: npa.ndarray) -> npa.ndarray:
"""Constructs the full kernel from its nonnegative quadrant.
Args:
@@ -83,24 +83,24 @@ def _quarter_to_full_kernel(arr: np.ndarray, pad_to: np.ndarray) -> np.ndarray:
Returns:
The complete kernel.
"""
pad_size = pad_to - 2 * np.array(arr.shape) + 1
pad_size = pad_to - 2 * npa.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 = npa.zeros((pad_size[0], arr.shape[1]))
bottom = npa.zeros((pad_size[0], arr.shape[1] - 1))
middle = npa.zeros((pad_to[0], pad_size[1]))
top_left = arr[:, :]
top_right = np.flipud(arr[1:, :])
bottom_left = np.fliplr(arr[:, 1:])
bottom_right = np.flipud(
np.fliplr(arr[1:, 1:])
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 np.concatenate(
return npa.concatenate(
(
np.concatenate((top_left, top, top_right)),
npa.concatenate((top_left, top, top_right)),
middle,
np.concatenate((bottom_left, bottom, bottom_right)),
npa.concatenate((bottom_left, bottom, bottom_right)),
),
axis=1,
)
@@ -113,34 +113,36 @@ def conic_filter(phase_mask,radius,space_grid):
phase_mask: the input phase mask to be filtered (2D)
radius: the radius of the conic filter in micrometers
space_grid: an object containing the size and resolution of the phase mask
Returns:
The filtered phase mask (2D)
'''
Nx = space_grid.get_Nx()
Ny = space_grid.get_Ny()
xv = np.arange(0, Nx//2-1) * (space_grid.get_resolution())
yv = np.arange(0, Ny//2-1) * (space_grid.get_resolution())
Xv, Yv = np.meshgrid(xv, yv)
h = np.where(Xv**2 + Yv**2 < radius**2, (1 - np.sqrt(abs(Xv**2 + Yv**2)) / radius),0)
#positive quadrant of the filter
xv = npa.arange(0, Nx//2+1) * (space_grid.get_resolution())
yv = npa.arange(0, Ny//2+1) * (space_grid.get_resolution())
#note that the filter's center is at (0,0)
Xv, Yv = npa.meshgrid(xv, yv)
h = _quarter_to_full_kernel(h, 3 * np.array([Nx, Ny]))
h = npa.where(Xv**2 + Yv**2 < radius**2, (1 - npa.sqrt(abs(Xv**2 + Yv**2)) / radius),0)
h = h / np.sum(h)
h = _quarter_to_full_kernel(h, 3 * npa.array([Nx, Ny]))
h = h / npa.sum(h)
#make a plot of the filter kernel
H = np.fft.fft2(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 = np.fft.fft2(phase_mask_padded)
phase_mask_fft = npa.fft.fft2(phase_mask_padded)
Y = phase_mask_fft * H
y = _centered(
np.real(np.fft.ifft2(Y)),
npa.real(npa.fft.ifft2(Y)),
(Nx, Ny)
)