mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 05:22:41 +02:00
520 lines
18 KiB
Python
520 lines
18 KiB
Python
from cdtools.tools import initializers
|
|
from cdtools.tools import propagators
|
|
from cdtools.tools import image_processing
|
|
|
|
import numpy as np
|
|
import torch as t
|
|
import pytest
|
|
import scipy.misc
|
|
from scipy import stats
|
|
from matplotlib import pyplot as plt
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def exit_waves_1():
|
|
# Import scipy test image and add a random phase
|
|
obj = scipy.misc.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])
|
|
#shape = t.Size([3,5])
|
|
wavelength = 1e-9
|
|
distance = 1#6e-3
|
|
rs_basis, full_shape, det_slice = \
|
|
initializers.exit_wave_geometry(basis, shape, wavelength,
|
|
distance, opt_for_fft=False)
|
|
|
|
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, full_shape, det_slice = \
|
|
initializers.exit_wave_geometry(basis, shape, wavelength,
|
|
distance, opt_for_fft=False)
|
|
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()
|
|
|
|
#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
|
|
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-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
|
|
|
|
|
|
Emz = np.conj(Ez)
|
|
|
|
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 it with the phase correction
|
|
asp = propagators.generate_angular_spectrum_propagator(
|
|
E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True,
|
|
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,remove_z_phase=True,
|
|
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)
|
|
|
|
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=False):
|
|
# 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))
|
|
if correct:
|
|
E = E * np.exp(-1j * k * Zs)
|
|
return E
|
|
|
|
|
|
# We make some rotation matrices to test
|
|
|
|
# This tests the straight ahead case
|
|
I = 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,I,Rinv, Rboth, Rboth,Rboth, Rall, Rall, Rall, I, 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 np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)):
|
|
#if True:
|
|
plt.close('all')
|
|
plt.imshow(np.angle(E0))
|
|
plt.figure()
|
|
plt.imshow(np.angle(Ez))
|
|
plt.figure()
|
|
plt.imshow(np.abs(Ez_t))
|
|
plt.figure()
|
|
plt.imshow(np.angle(Ez_t))
|
|
plt.figure()
|
|
plt.imshow(np.abs(Ez-Ez_t)/np.max(np.abs(Ez)))
|
|
plt.show()
|
|
|
|
assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
|
|
|
|
|
|
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy()
|
|
|
|
# Again, 10^-3 is about all the accuracy we can expect
|
|
assert np.max(np.abs(E0-Em0_t)) < 1e-3 * np.max(np.abs(E0))
|
|
|
|
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))
|