mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-10-12 11:00:26 +02:00
linting test_plotting.py and test_propagators.py
This commit is contained in:
gnzng
@@ -1,47 +1,50 @@
|
||||
from cdtools.tools import plotting
|
||||
from cdtools.tools import initializers
|
||||
import numpy as np
|
||||
import torch as t
|
||||
import scipy.datasets
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from cdtools.tools import plotting
|
||||
from cdtools.tools import initializers
|
||||
|
||||
|
||||
def test_plot_amplitude(show_plot):
|
||||
# Test with tensor
|
||||
im = t.as_tensor(scipy.datasets.ascent(),dtype=t.complex128)
|
||||
plotting.plot_amplitude(im, basis = np.array([[0,-1], [-1,0], [0,0]]), title = 'Test Amplitude')
|
||||
im = t.as_tensor(scipy.datasets.ascent(), dtype=t.complex128)
|
||||
plotting.plot_amplitude(im, basis=np.array([[0, -1], [-1, 0], [0, 0]]), title='Test Amplitude')
|
||||
if show_plot:
|
||||
plt.show()
|
||||
|
||||
# Test with numpy array
|
||||
im = scipy.datasets.ascent().astype(np.complex128)
|
||||
plotting.plot_amplitude(im, title = 'Test Amplitude')
|
||||
plotting.plot_amplitude(im, title='Test Amplitude')
|
||||
if show_plot:
|
||||
plt.show()
|
||||
|
||||
|
||||
def test_plot_phase(show_plot):
|
||||
# Test with tensor
|
||||
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
|
||||
plotting.plot_phase(im, title = 'Test Phase')
|
||||
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
|
||||
plotting.plot_phase(im, title='Test Phase')
|
||||
if show_plot:
|
||||
plt.show()
|
||||
|
||||
# Test with numpy array
|
||||
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy()
|
||||
plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[0,-1], [-1,0], [0,0]]))
|
||||
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1]).numpy()
|
||||
plotting.plot_phase(im, title='Test Phase', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
|
||||
if show_plot:
|
||||
plt.show()
|
||||
|
||||
|
||||
def test_plot_colorized(show_plot):
|
||||
# Test with tensor
|
||||
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
|
||||
gaussian = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
|
||||
im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64)
|
||||
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[0,-1], [-1,0], [0,0]]))
|
||||
plotting.plot_colorized(im, title='Test Colorize', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
|
||||
if show_plot:
|
||||
plt.show()
|
||||
|
||||
# Test with numpy array
|
||||
im = im.numpy()
|
||||
plotting.plot_colorized(im, title = 'Test Colorize')
|
||||
plotting.plot_colorized(im, title='Test Colorize')
|
||||
if show_plot:
|
||||
plt.show()
|
||||
|
||||
+194
-211
@@ -1,7 +1,3 @@
|
||||
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
|
||||
@@ -9,35 +5,37 @@ 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 = 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()))
|
||||
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')))
|
||||
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)))
|
||||
assert (np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
|
||||
|
||||
|
||||
def test_generate_high_NA_k_intensity_map():
|
||||
@@ -45,9 +43,9 @@ 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])
|
||||
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 = \
|
||||
@@ -60,7 +58,7 @@ def test_generate_high_NA_k_intensity_map():
|
||||
# generate a good test exit wave
|
||||
i = (np.arange(478) - 240)
|
||||
j = (np.arange(573) - 270)
|
||||
Is,Js = np.meshgrid(i,j,indexing='ij')
|
||||
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)
|
||||
|
||||
@@ -74,37 +72,36 @@ def test_generate_high_NA_k_intensity_map():
|
||||
# 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
|
||||
# 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)
|
||||
# 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]
|
||||
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()
|
||||
|
||||
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),
|
||||
@@ -114,8 +111,9 @@ def test_generate_high_NA_k_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
|
||||
|
||||
# assert 0
|
||||
|
||||
|
||||
def test_near_field_direction():
|
||||
#
|
||||
@@ -135,19 +133,19 @@ def test_near_field_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)
|
||||
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
|
||||
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)
|
||||
E0_real.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex64)
|
||||
|
||||
Ez_real = propagators.near_field(E0_real,asp)
|
||||
Ez_real = propagators.near_field(E0_real, asp)
|
||||
|
||||
centroid = image_processing.centroid(t.abs(Ez_real))
|
||||
|
||||
@@ -155,35 +153,35 @@ def test_near_field_direction():
|
||||
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()
|
||||
|
||||
|
||||
# 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)
|
||||
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
|
||||
wavelength = 3e-9 # nm
|
||||
sigma = 20e-9 # nm
|
||||
z = 1000e-9 # nm
|
||||
|
||||
k = 2 * np.pi / wavelength
|
||||
w0 = np.sqrt(2)*sigma
|
||||
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)
|
||||
|
||||
Rz = z * (1 + (zr / z)**2)
|
||||
|
||||
E0 = np.exp(-Rs**2 / w0**2)
|
||||
|
||||
# The analytical expression for propagation of a gaussian beam in the
|
||||
@@ -208,52 +206,49 @@ def test_near_field():
|
||||
# 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 = 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)
|
||||
E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
|
||||
|
||||
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
|
||||
|
||||
|
||||
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))
|
||||
|
||||
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()
|
||||
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))
|
||||
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,
|
||||
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
|
||||
|
||||
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)
|
||||
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()
|
||||
t.real(asp[10, 10]).backward()
|
||||
assert z.grad != 0
|
||||
assert spacing.grad[0] != 0
|
||||
assert wavelength.grad !=0
|
||||
|
||||
assert wavelength.grad != 0
|
||||
|
||||
|
||||
def test_generalized_near_field():
|
||||
|
||||
@@ -264,26 +259,25 @@ def test_generalized_near_field():
|
||||
# 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)
|
||||
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])
|
||||
|
||||
Positions = np.stack([Xs_0,Ys_0,Zs_0])
|
||||
|
||||
# assert 0
|
||||
wavelength = 3e-9 #nm
|
||||
sigma = 20e-9 #nm
|
||||
z = 1000e-9 #nm
|
||||
wavelength = 3e-9 # nm
|
||||
sigma = 20e-9 # nm
|
||||
z = 1000e-9 # nm
|
||||
|
||||
k = 2 * np.pi / wavelength
|
||||
w0 = np.sqrt(2)*sigma
|
||||
w0 = np.sqrt(2) * sigma
|
||||
zr = np.pi * w0**2 / wavelength
|
||||
|
||||
|
||||
# The analytical expression for propagation of a gaussian beam in the
|
||||
# paraxial approx
|
||||
@@ -292,140 +286,132 @@ def test_generalized_near_field():
|
||||
|
||||
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))
|
||||
|
||||
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:
|
||||
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)))
|
||||
|
||||
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
|
||||
I = np.eye(3)
|
||||
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)]])
|
||||
|
||||
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)
|
||||
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]])
|
||||
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]])
|
||||
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))
|
||||
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])
|
||||
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])
|
||||
|
||||
z_dir_large = np.array([0, 0, 10])
|
||||
|
||||
# And finally some offset vectors
|
||||
|
||||
# This checks straight ahead
|
||||
z_offset = np.array([0,0,z])
|
||||
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])
|
||||
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])
|
||||
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,
|
||||
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],
|
||||
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):
|
||||
|
||||
# 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)
|
||||
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]
|
||||
Xs_prop, Ys_prop, Zs_prop = np.stack([Xs, Ys, Zs_0]) \
|
||||
+ offset_vec[:, None, None]
|
||||
|
||||
print('Propagate Along', propagation_vec)
|
||||
|
||||
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)
|
||||
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)
|
||||
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)
|
||||
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,
|
||||
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)
|
||||
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)
|
||||
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()
|
||||
|
||||
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:
|
||||
# if True:
|
||||
plt.close('all')
|
||||
plt.figure()
|
||||
plt.imshow(np.angle(E0))
|
||||
@@ -443,70 +429,67 @@ def test_generalized_near_field():
|
||||
plt.imshow(np.angle(Ez_t))
|
||||
plt.title('Angle of numerically calculated Ez')
|
||||
plt.figure()
|
||||
plt.imshow(np.abs(Ez-Ez_t))#/np.max(np.abs(Ez)))
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
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)
|
||||
|
||||
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()
|
||||
|
||||
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)
|
||||
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()
|
||||
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))
|
||||
|
||||
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
|
||||
Ys, Xs = np.meshgrid(y, x)
|
||||
Rs = np.sqrt(Xs**2 + Ys**2)
|
||||
|
||||
w0 = np.sqrt(2)*sigma
|
||||
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.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)
|
||||
|
||||
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))
|
||||
assert t.all(t.isclose(E0, E_backprop))
|
||||
|
||||
Reference in New Issue
Block a user