odak.learn.wave.custom¶
A definition to calculate convolution based Fresnel approximation for beam propagation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
field |
torch.complex |
Complex field (MxN). |
required |
kernel |
torch.complex |
Custom complex kernel for beam propagation. |
required |
Returns:
| Type | Description |
|---|---|
torch.complex |
Final complex field (MxN). |
Source code in odak/learn/wave/classical.py
def custom(field, kernel, zero_padding = False):
"""
A definition to calculate convolution based Fresnel approximation for beam propagation.
Parameters
----------
field : torch.complex
Complex field (MxN).
kernel : torch.complex
Custom complex kernel for beam propagation.
Returns
-------
result : torch.complex
Final complex field (MxN).
"""
if type(kernel) == type(None):
H = torch.zeros(field.shape).to(field.device)
else:
H = kernel
U1 = torch.fft.fftshift(torch.fft.fft2(torch.fft.fftshift(field)))
if zero_padding == False:
U2 = H*U1
elif zero_padding == True:
U2 = zero_pad(H*U1)
result = torch.fft.ifftshift(torch.fft.ifft2(torch.fft.ifftshift(U2)))
return result
Notes¶
Unless you know what you are doing, we do not suggest you to use this function directly. Rather stick to odak.learn.wave.propagate_beam for your beam propagation code. Note that this function can also be used as convolution operation between two complex arrays.