cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-20 01:22:09 +02:00
Code Issues Packages Projects Releases Wiki Activity

Debug and test the off-axis propagation code

Browse Source
This commit is contained in:
Abe Levitan
2020-04-24 14:16:44 -04:00
parent 071eebd37b
commit d586ed914e
2 changed files with 144 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/tools/propagators.py
+33 -7
View File
@@ -72,7 +72,7 @@ def inverse_far_field(wavefront):
return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, **kwargs):
def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, **kwargs):
"""Generates an angular-spectrum based near-field propagator from experimental quantities
This function generates an angular-spectrum based near field
@@ -96,6 +96,8 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, *
The wavelength of light to simulate propagation of
z : float
The distance to simulate propagation over
remove_z_phase : bool
Default False, whether to remove the dominant z-direction phase dependence
Returns
-------
@@ -113,6 +115,9 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, *
propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z)
if remove_z_phase:
propagator *= np.exp(-1j * k0 * z)
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
@@ -139,12 +144,15 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
is designed to work on any wavefield defined on an array of
parallelograms.
In addition, if propagate_along_offset is true, there is an assumed phase
In addition, if propagate_along_offset is True, there is an assumed phase
ramp applied to the wavefield before propagation, defined such that a
feature with uniform phase will propagate along the direction of the
defined offset vector. This decision provides the best numerical
stability and allows for the simple setup of light fields copropagating
with the coordinate system.
defined offset vector. This will also remove the phase variation along
the propagation direction, because it makes the most physical sense to
regard this choice as removing the dominant phase variation in 3D, allowing
for the generation of a smoothly varying wavefield over 3D volumes.
This decision provides the best numerical stability and allows for the
simple setup of light fields copropagating with the coordinate system.
Parameters
@@ -157,6 +165,8 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
The wavelength of light to simulate propagation of
propagation_vector : array
The displacement to propagate the wavefield along.
propagate_along_offset : bool
Optional, whether to include an implied phase ramp to propagate uniform phase features along the offset direction
Returns
-------
@@ -186,7 +196,6 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
# This may have a sign error - must be checked
phase_mask = np.exp(1j * np.tensordot(offset_vector,K_xyz,axes=1))
# Next, we apply a shift to the k-space vectors which sets up
# propagation such that a uniform phase object will propagate along the
@@ -198,10 +207,22 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
perpendicular_dir /= np.linalg.norm(perpendicular_dir)
offset_perpendicular = np.dot(perpendicular_dir, offset_vector)
offset_parallel = offset_vector - perpendicular_dir * offset_perpendicular
print(offset_perpendicular)
print(offset_parallel)
k0 = 2*np.pi/wavelength
k_offset = offset_parallel * k0 / np.sqrt(offset_perpendicular**2 +
np.linalg.norm(offset_parallel)**2)
K_xyz = K_xyz - k_offset[:,None,None]
# Only implement the shift if the flag is set to True
if propagate_along_offset:
K_xyz = K_xyz + k_offset[:,None,None]
# we also need to remove the z-dependence on the phase
# This time, though, the z-dependence actually has to do with
# the oput of plane component of k at the central offset. Normally
# this is 0, so the z-component is just k0, but not in this case
phase_mask *= np.exp(-1j * np.sqrt(k0**2 - np.linalg.norm(k_offset)**2)
* offset_perpendicular)
# Redefine this as complex so the square root properly gives
# k>k0 components imaginary frequencies
@@ -212,6 +233,11 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
* offset_perpendicular)
propagator *= phase_mask
# This removes the z-dependence on the phase:
#if propagate_along_offset:
# print(offset_perpendicular)
# propagator *= np.exp(-1j * k0 * offset_perpendicular)
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
tests/tools/test_propagators.py
+111 -22
View File
@@ -9,7 +9,7 @@ import torch as t
import pytest
import scipy.misc
from scipy.fftpack import fftshift, ifftshift
from matplotlib import pyplot as plt
@pytest.fixture(scope='module')
@@ -68,9 +68,13 @@ def test_near_field():
# The analytical expression for propagation of a gaussian beam in the
# paraxial approx
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.float64)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
@@ -87,52 +91,137 @@ def test_near_field():
# 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.float64)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
# 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(cmath.complex_to_torch(E0),asp)
Emz_t = cmath.torch_to_complex(Emz_t)
# 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))
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]])
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)
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
propagation_vector = np.array([0,0,z])
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
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * ( z + Rs**2 / (2 * Rz)) + 1j * np.arctan(z / zr))
def get_w(Zs):
return w0 * np.sqrt(1 + (Zs / zr)**2)
basis= np.array([[0,-1e-9],[-1.5e-9,0],[0,0]])
propagation_vector = np.array([0,0,z])
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,basis,wavelength,propagation_vector,dtype=t.float64)
#assert False
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
def get_inv_R(Zs):
return Zs / (Zs**2 + zr**2)
# 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))
def get_E(Xs, Ys, Zs, correct=False):
# 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
# 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 a shearing and a rotation together
Rall = np.matmul(Rboth,Rshear)
rot_mats = [I,Ry, Rboth, Rshear, Rboth]
purposes = ['standard','y-rot','both-rot','shear','shear-rot']
for purpose,rot_mat in zip(purposes,rot_mats):
print('Testing', purpose)
Xs,Ys,Zs_0 = np.tensordot(rot_mat,Positions,axes=1)
new_basis = np.dot(rot_mat, basis)
Zs_prop = Zs_0 + z
# Check that it works both with the explicit and implicit phase ramps
for prop_oo in [False, True]:
print('Propagate Along Offset =',prop_oo)
E0 = get_E(Xs,Ys,Zs_0, correct=prop_oo)
Ez = get_E(Xs,Ys,Zs_prop, correct=prop_oo)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,propagation_vector,
dtype=t.float64, propagate_along_offset=prop_oo)
Emz = np.conj(Ez)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp)
Emz_t = cmath.torch_to_complex(Emz_t)
# 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))
Em0_t = propagators.inverse_near_field(cmath.complex_to_torch(Ez),asp)
Em0_t = cmath.torch_to_complex(Em0_t)
# 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')
# 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))
def test_inverse_near_field():