mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 21:42:39 +02:00
496 lines
18 KiB
Python
496 lines
18 KiB
Python
import numpy as np
|
|
import torch as t
|
|
import pytest
|
|
import scipy.datasets
|
|
from scipy import stats
|
|
from matplotlib import pyplot as plt
|
|
|
|
from cdtools.tools import initializers
|
|
from cdtools.tools import propagators
|
|
from cdtools.tools import image_processing
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def exit_waves_1():
|
|
# Import scipy test image and add a random phase
|
|
obj = scipy.datasets.ascent()[0:64, 0:64].astype(np.complex128)
|
|
arr = np.random.random_sample((64, 64))
|
|
obj *= (arr + (1 - arr**2)**.5 * 1j)
|
|
obj = t.as_tensor(obj)
|
|
|
|
# Construct wavefront from image
|
|
probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
|
|
return probe * obj
|
|
|
|
|
|
def test_far_field(exit_waves_1):
|
|
# Far field diffraction patterns calculated by numpy with zero frequency in center
|
|
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho'))
|
|
|
|
assert (np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
|
|
|
|
|
|
def test_inverse_far_field(exit_waves_1):
|
|
# We want the inverse far field to map back to the exit waves with no intensity corrections
|
|
# Far field result for exit waves calculated with numpy
|
|
far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho')))
|
|
|
|
assert (np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
|
|
|
|
|
|
def test_generate_high_NA_k_intensity_map():
|
|
|
|
# We need to generate a plausible scenario. I will start
|
|
# by using the initializer to generate a reasonable exit wave geometry
|
|
# and detector pair
|
|
basis = t.Tensor([[0, -30e-6, 0],
|
|
[-20e-6, 0, 0]]).transpose(0, 1)
|
|
shape = t.Size([478, 573])
|
|
wavelength = 1e-9
|
|
distance = 1
|
|
rs_basis = \
|
|
initializers.exit_wave_geometry(basis, shape, wavelength, distance)
|
|
|
|
k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
|
|
rs_basis, basis, shape, distance, wavelength,
|
|
dtype=t.float32)
|
|
|
|
# generate a good test exit wave
|
|
i = (np.arange(478) - 240)
|
|
j = (np.arange(573) - 270)
|
|
Is, Js = np.meshgrid(i, j, indexing='ij')
|
|
wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128)
|
|
t_wavefield = t.as_tensor(wavefield, dtype=t.complex64)
|
|
|
|
high_NA_propagated = propagators.high_NA_far_field(
|
|
t_wavefield, k_map, intensity_map=intensity_map)
|
|
low_NA_propagated = propagators.far_field(t_wavefield)
|
|
|
|
low_NA = low_NA_propagated.numpy()
|
|
high_NA = high_NA_propagated.numpy()
|
|
|
|
# Checking first that for a low-NA propagation they give the same result
|
|
# 1e-4 tolerance seems to be reasonable in this comparison given my
|
|
# exploration with the code
|
|
# assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
|
|
|
|
# Now I will explore some results with a tilted sample
|
|
# print(rs_basis)
|
|
# print(rs_basis_tilted)
|
|
distance = 0.01 # 6e-3
|
|
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
|
|
rs_basis_tilted = rs_basis.clone()
|
|
rs_basis_tilted[2, 1] = rs_basis_tilted[0, 1]
|
|
|
|
k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
|
|
rs_basis_tilted, basis, shape, distance, wavelength,
|
|
dtype=t.float32)
|
|
|
|
high_NA_propagated = propagators.high_NA_far_field(
|
|
t_wavefield, k_map, intensity_map=intensity_map)
|
|
low_NA_propagated = propagators.far_field(t_wavefield)
|
|
|
|
low_NA = low_NA_propagated.numpy()
|
|
high_NA = high_NA_propagated.numpy()
|
|
|
|
print(low_NA.shape, high_NA.shape)
|
|
|
|
# plt.close('all')
|
|
# plt.imshow(np.abs(low_NA))
|
|
# plt.figure()
|
|
# plt.imshow(np.abs(high_NA))
|
|
# plt.colorbar()
|
|
# plt.imshow(np.abs(wavefield))
|
|
# plt.show()
|
|
|
|
# Now I want to test that it doesn't crash for wavefields of various shapes
|
|
propagators.high_NA_far_field(t_wavefield.unsqueeze(0),
|
|
k_map, intensity_map=intensity_map)
|
|
propagators.high_NA_far_field(t_wavefield.unsqueeze(0).unsqueeze(0),
|
|
k_map, intensity_map=intensity_map)
|
|
|
|
# I believe this works, but I still would like to get a second method for
|
|
# simulating at least one diffraction pattern as an independent check
|
|
|
|
# assert 0
|
|
|
|
|
|
def test_near_field_direction():
|
|
#
|
|
# The idea is to make sure that near field propagation is consistent
|
|
# with far-field propagation. Because we consider the Fourier
|
|
# transform to be the far-field propagator, we should find that
|
|
# near-field propagation becomes far-field propagation in the limit
|
|
# that z goes to positive infinity. This consistency requirement
|
|
# forces a specific choice of sign for near-field propagation - choosing
|
|
# it incorrectly will make the correct limit obtain at negative
|
|
# infinity.
|
|
#
|
|
# We can check consistency by setting up a gaussian wavefield with an
|
|
# offset in the far field. Then, we can do some near-field propagation,
|
|
# and we will ensure that the centroid of the beam moves in the correct
|
|
# direction. If the sign is incorrect, it will move in the wrong direction.
|
|
|
|
x = (np.arange(901) - 400)
|
|
y = (np.arange(1200) - 500)
|
|
Ys, Xs = np.meshgrid(y, x)
|
|
Rs = np.sqrt(Xs**2 + Ys**2)
|
|
|
|
E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)), dtype=t.complex64)
|
|
E0_real = propagators.inverse_far_field(E0_fourier)
|
|
# This is in the top-left corner in Fourier space
|
|
|
|
wavelength = 3e-9 # nm
|
|
z = 1000e-9 # nm
|
|
asp = propagators.generate_angular_spectrum_propagator(
|
|
E0_real.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex64)
|
|
|
|
Ez_real = propagators.near_field(E0_real, asp)
|
|
|
|
centroid = image_processing.centroid(t.abs(Ez_real))
|
|
|
|
# Assert it's in top half
|
|
assert centroid[0] < Ez_real.shape[0] // 2
|
|
# Assert it's in left half
|
|
assert centroid[1] < Ez_real.shape[1] // 2
|
|
|
|
# plt.imshow(t.abs(E0_fourier))
|
|
# plt.figure()
|
|
# plt.imshow(t.abs(E0_real))
|
|
# plt.figure()
|
|
# plt.imshow(t.abs(Ez_real))
|
|
# plt.show()
|
|
|
|
|
|
def test_near_field():
|
|
|
|
# The strategy is to compare the propagation of a gaussian beam to
|
|
# the propagation in the paraxial approximation.
|
|
|
|
x = (np.arange(901) - 450) * 1.5e-9
|
|
y = (np.arange(1200) - 600) * 1e-9
|
|
Ys, Xs = np.meshgrid(y, x)
|
|
Rs = np.sqrt(Xs**2 + Ys**2)
|
|
|
|
wavelength = 3e-9 # nm
|
|
sigma = 20e-9 # nm
|
|
z = 1000e-9 # nm
|
|
|
|
k = 2 * np.pi / wavelength
|
|
w0 = np.sqrt(2) * sigma
|
|
zr = np.pi * w0**2 / wavelength
|
|
wz = w0 * np.sqrt(1 + (z / zr)**2)
|
|
Rz = z * (1 + (zr / z)**2)
|
|
|
|
E0 = np.exp(-Rs**2 / w0**2)
|
|
|
|
# The analytical expression for propagation of a gaussian beam in the
|
|
# paraxial approx
|
|
# The sign convention I'm using here is opposite from what is on
|
|
# wikipedia (as of June 2021), but it is I believe the most sensible
|
|
# choice. By choosing e^(ikx) to be the plane wave that's propagating
|
|
# in the x direction, we preserve two nice properties:
|
|
# 1) The time-dependence is e^(-iwt+ikx), which meshes with the choice
|
|
# we like to make as physicists of e^(-iwt) for time-dependence
|
|
# 2) This makes it so that the Fourier transform does far-field
|
|
# propagation in direction of light propagation. The logic is as
|
|
# follows:
|
|
# a) We must assume that light only propagates one direction through
|
|
# the plane, in order to have well-defined propagation
|
|
# b) Fixing that direction to be the positive direciton of propagation,
|
|
# we can consider the in-plane image of a plane wave propagating
|
|
# through the plane of interest.
|
|
# c) We can ask whether the 2D Fourier transform of that in-plane image
|
|
# produces a spot on the same side of zero as the in-plane component
|
|
# of k.
|
|
# If we choose e^(-ikx) to represent light propagating along K, the
|
|
# answer is no, and we find we have to use the inverse FT instead.
|
|
# Thus, e^(ikx) is the right choice here.
|
|
|
|
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * (z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
|
|
Ez_nozphase = Ez * np.exp(-1j * k * z)
|
|
|
|
# First we check it normally
|
|
asp = propagators.generate_angular_spectrum_propagator(
|
|
E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
|
|
|
|
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
|
|
|
|
# Check for at least 10^-3 relative accuracy in this scenario
|
|
assert np.max(np.abs(Ez_nozphase - Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
|
|
|
|
Emz = np.conj(Ez_nozphase)
|
|
|
|
Emz_t = propagators.inverse_near_field(t.as_tensor(E0), asp).numpy()
|
|
|
|
# Again, 10^-3 is about all the accuracy we can expect
|
|
assert np.max(np.abs(Emz - Emz_t)) < 1e-3 * np.max(np.abs(Emz))
|
|
|
|
# Then, we check that the bandlimiting at least does something
|
|
asp = propagators.generate_angular_spectrum_propagator(
|
|
E0.shape, (1.5e-9, 1e-9), wavelength, z,
|
|
dtype=t.complex128, bandlimit=0.3)
|
|
|
|
assert asp[140, 0] == 0
|
|
assert asp[0, 180] == 0
|
|
assert asp[130, 0] != 0
|
|
assert asp[0, 175] != 0
|
|
|
|
# Then, we check that automatic differentiation works
|
|
z = t.tensor([z], requires_grad=True)
|
|
spacing = t.tensor((1.5e-9, 1e-9), requires_grad=True)
|
|
wavelength = t.tensor([wavelength], requires_grad=True)
|
|
|
|
asp = propagators.generate_angular_spectrum_propagator(
|
|
E0.shape, spacing, wavelength, z)
|
|
|
|
t.real(asp[10, 10]).backward()
|
|
assert z.grad != 0
|
|
assert spacing.grad[0] != 0
|
|
assert wavelength.grad != 0
|
|
|
|
|
|
def test_generalized_near_field():
|
|
|
|
# The strategy is to compare the propagation of a gaussian beam to
|
|
# the propagation in the paraxial approximation.
|
|
|
|
# For this one, we want to test it on a rotated coordinate system
|
|
# First, we should do a test with the phase ramp along the z direction
|
|
# explicitly included
|
|
|
|
basis = np.array([[0, -1.5e-9], [-1e-9, 0], [0, 0]])
|
|
i_vec, j_vec = np.arange(901) - 450, np.arange(1200) - 600
|
|
Is, Js = np.meshgrid(i_vec, j_vec, indexing='ij')
|
|
Xs_0, Ys_0, Zs_0 = np.tensordot(basis, np.stack([Is, Js]), axes=1)
|
|
# x = (np.arange(901) - 450) * 1.5e-9
|
|
# y = (np.arange(1200) - 600) * 1e-9
|
|
# Xs_0,Ys_0 = np.meshgrid(x,y)
|
|
# Zs_0 = np.zeros(Xs_0.shape)
|
|
|
|
Positions = np.stack([Xs_0, Ys_0, Zs_0])
|
|
|
|
# assert 0
|
|
wavelength = 3e-9 # nm
|
|
sigma = 20e-9 # nm
|
|
z = 1000e-9 # nm
|
|
|
|
k = 2 * np.pi / wavelength
|
|
w0 = np.sqrt(2) * sigma
|
|
zr = np.pi * w0**2 / wavelength
|
|
|
|
# The analytical expression for propagation of a gaussian beam in the
|
|
# paraxial approx
|
|
def get_w(Zs):
|
|
return w0 * np.sqrt(1 + (Zs / zr)**2)
|
|
|
|
def get_inv_R(Zs):
|
|
return Zs / (Zs**2 + zr**2)
|
|
|
|
def get_E(Xs, Ys, Zs, correct=True):
|
|
# Again, this follows the convention opposite from Wikipedia. See
|
|
# the note in test_angular_spectrum_propagator.
|
|
|
|
# if correct is True, remove the e^(ikz) dependence
|
|
|
|
Rs_sq = Xs**2 + Ys**2
|
|
Wzs = get_w(Zs)
|
|
E = w0 / Wzs * np.exp(-Rs_sq / Wzs**2) *\
|
|
np.exp(1j * k * (Zs + Rs_sq * get_inv_R(Zs) / 2) - 1j * np.arctan(Zs / zr))
|
|
|
|
# This removes the z-dependence of the phase
|
|
if correct:
|
|
E = E * np.exp(-1j * k * Zs)
|
|
return E
|
|
|
|
def check_equiv(analytical, numerical):
|
|
phase = np.angle(np.mean(numerical.conj() * analytical))
|
|
comp = np.exp(1j * phase) * numerical
|
|
return (np.max(np.abs(analytical - comp)) < 1e-3 * np.max(np.abs(analytical)))
|
|
|
|
# We make some rotation matrices to test
|
|
|
|
# This tests the straight ahead case
|
|
IdentityMatrix = np.eye(3)
|
|
|
|
# This tests a rotation about the y axis
|
|
th = np.deg2rad(5)
|
|
Ry = np.array([[np.cos(th), 0, np.sin(th)],
|
|
[0, 1, 0],
|
|
[-np.sin(th), 0, np.cos(th)]])
|
|
|
|
# This tests a rotation about two axes
|
|
phi = np.deg2rad(2)
|
|
Rx = np.array([[1, 0, 0],
|
|
[0, np.cos(phi), -np.sin(phi)],
|
|
[0, np.sin(phi), np.cos(phi)]])
|
|
Rboth = np.matmul(Rx, Ry)
|
|
|
|
# This tests a shearing
|
|
shear = 0.23
|
|
Rshear = np.array([[1, shear, 0],
|
|
[0, 1, 0],
|
|
[0, 0, 1]])
|
|
|
|
# This tests an inversion of the axes
|
|
Rinv = np.array([[-1, 0, 0],
|
|
[0, -1, 0],
|
|
[0, 0, -1]])
|
|
|
|
# This tests a reflection about the y-z plane
|
|
Rrefl = np.array([[-1, 0, 0],
|
|
[0, 1, 0],
|
|
[0, 0, -1]])
|
|
|
|
# This tests a shearing and a rotation together
|
|
Rall = np.matmul(Rrefl, np.matmul(Rboth, Rshear))
|
|
|
|
# And we make some propagation vectors to test:
|
|
|
|
# This is along the z direction
|
|
z_dir = np.array([0, 0, 1])
|
|
|
|
# This checks that it's not sensitive to the magnitude
|
|
z_dir_large = np.array([0, 0, 10])
|
|
|
|
# And finally some offset vectors
|
|
|
|
# This checks straight ahead
|
|
z_offset = np.array([0, 0, z])
|
|
|
|
# This checks with an offset in x and y
|
|
shear_offset = np.array([0.1 * z, -0.03 * z, z])
|
|
|
|
# This checks with an offset in x and y, with negative z
|
|
shear_back_offset = np.array([0.1 * z, -0.03 * z, -z])
|
|
|
|
rot_mats = [Rrefl, IdentityMatrix, Rinv, Rboth, Rboth, Rboth, Rall, Rall, Rall, IdentityMatrix, Rall]
|
|
offset_vecs = [z_offset] * 8 + [shear_offset] + [shear_back_offset] * 2
|
|
propagation_vecs = ['perp', 'offset', z_dir,
|
|
'perp', 'offset', z_dir_large,
|
|
'perp', 'offset', z_dir_large,
|
|
z_dir, z_dir_large]
|
|
purposes = ['standard'] * 3 + ['both-rot'] * 3 + ['shear-rot'] * 3 + ['backward'] * 2
|
|
|
|
# rot_mats = [Ry.transpose()]
|
|
# offset = np.cross(np.dot(Ry.transpose(),basis)[:,0],
|
|
# np.dot(Ry.transpose(),basis)[:,1])
|
|
# offset /= np.linalg.norm(offset) / 3e-6
|
|
# offset_vecs = [-offset]#[shear_offset]
|
|
# propagation_vecs = [z_dir]
|
|
# purposes=['meh']
|
|
|
|
for purpose, rot_mat, offset_vec, propagation_vec in zip(purposes, rot_mats, offset_vecs, propagation_vecs):
|
|
print('Testing', purpose)
|
|
Xs, Ys, Zs_0 = np.tensordot(rot_mat, Positions, axes=1)
|
|
new_basis = np.dot(rot_mat, basis)
|
|
Xs_prop, Ys_prop, Zs_prop = np.stack([Xs, Ys, Zs_0]) \
|
|
+ offset_vec[:, None, None]
|
|
|
|
print('Propagate Along', propagation_vec)
|
|
|
|
if str(propagation_vec) == 'perp':
|
|
E0 = get_E(Xs, Ys, Zs_0, correct=False)
|
|
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=False)
|
|
asp = propagators.generate_generalized_angular_spectrum_propagator(
|
|
E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128)
|
|
elif str(propagation_vec) == 'offset':
|
|
E0 = get_E(Xs, Ys, Zs_0, correct=True)
|
|
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
|
|
asp = propagators.generate_generalized_angular_spectrum_propagator(
|
|
E0.shape, new_basis, wavelength, offset_vec,
|
|
dtype=t.complex128, propagate_along_offset=True)
|
|
else:
|
|
E0 = get_E(Xs, Ys, Zs_0, correct=True)
|
|
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
|
|
asp = propagators.generate_generalized_angular_spectrum_propagator(
|
|
E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128,
|
|
propagation_vector=propagation_vec)
|
|
|
|
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
|
|
|
|
# Check for at least 10^-3 relative accuracy in this scenario
|
|
if not check_equiv(Ez, Ez_t):
|
|
# if True:
|
|
plt.close('all')
|
|
plt.figure()
|
|
plt.imshow(np.angle(E0))
|
|
plt.title('Angle of E0')
|
|
plt.figure()
|
|
plt.imshow(np.angle(Ez))
|
|
plt.title('Angle of analytically calculated Ez')
|
|
plt.figure()
|
|
plt.imshow(np.abs(Ez))
|
|
plt.title('Magnitude of analytically calculated Ez')
|
|
plt.figure()
|
|
plt.imshow(np.abs(Ez_t))
|
|
plt.title('Magnitude of numerically calculated Ez')
|
|
plt.figure()
|
|
plt.imshow(np.angle(Ez_t))
|
|
plt.title('Angle of numerically calculated Ez')
|
|
plt.figure()
|
|
plt.imshow(np.abs(Ez - Ez_t))
|
|
plt.title('Magnitude of difference')
|
|
plt.show()
|
|
|
|
assert check_equiv(Ez, Ez_t)
|
|
|
|
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez), asp).numpy()
|
|
|
|
assert check_equiv(E0, Em0_t)
|
|
|
|
print('Test Successful')
|
|
|
|
# One final test, to see if any of a few arbitrary rotations will
|
|
# change the predicted propagation if everything else is kept
|
|
# constant
|
|
Rrands = [stats.ortho_group.rvs(3) for i in range(3)]
|
|
Xs, Ys, Zs_0 = np.tensordot(Rboth, Positions, axes=1)
|
|
new_basis = np.dot(rot_mat, basis)
|
|
offset_vec = shear_back_offset
|
|
propagation_vec = z_dir
|
|
|
|
E0 = get_E(Xs, Ys, Zs_0, correct=True)
|
|
asp = propagators.generate_generalized_angular_spectrum_propagator(
|
|
E0.shape, new_basis, wavelength, offset_vec,
|
|
dtype=t.complex128, propagation_vector=propagation_vec)
|
|
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
|
|
|
|
for Rrand in Rrands:
|
|
Xs, Ys, Zs_0 = np.tensordot(Rrand, np.tensordot(Rboth, Positions, axes=1), axes=1)
|
|
rot_offset = np.dot(Rrand, offset_vec)
|
|
rot_basis = np.dot(Rrand, new_basis)
|
|
rot_prop = np.dot(Rrand, propagation_vec)
|
|
asp = propagators.generate_generalized_angular_spectrum_propagator(
|
|
E0.shape, rot_basis, wavelength, rot_offset,
|
|
dtype=t.complex128, propagation_vector=rot_prop)
|
|
Ez_rot_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
|
|
|
|
assert np.max(np.abs(Ez_t - Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
|
|
|
|
|
|
def test_inverse_near_field():
|
|
|
|
x = (np.arange(800) - 400) * 1.5e-9
|
|
y = (np.arange(1200) - 600) * 1e-9
|
|
Ys, Xs = np.meshgrid(y, x)
|
|
Rs = np.sqrt(Xs**2 + Ys**2)
|
|
|
|
wavelength = 3e-9 # nm
|
|
sigma = 20e-9 # nm
|
|
z = 1000e-9 # nm
|
|
|
|
w0 = np.sqrt(2) * sigma
|
|
E0 = np.exp(-Rs**2 / w0**2)
|
|
|
|
asp = propagators.generate_angular_spectrum_propagator(
|
|
E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
|
|
|
|
E0 = t.as_tensor(E0, dtype=t.complex128)
|
|
E_prop = propagators.near_field(E0, asp)
|
|
|
|
E_backprop = propagators.inverse_near_field(E_prop, asp)
|
|
|
|
# We just want to check that it actually is the inverse
|
|
assert t.all(t.isclose(E0, E_backprop))
|