mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 05:22:41 +02:00
115 lines
3.5 KiB
Python
115 lines
3.5 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
|
|
from CDTools.tools import cmath
|
|
from CDTools.tools import initializers
|
|
from CDTools.tools import propagators
|
|
|
|
import numpy as np
|
|
import torch as t
|
|
import pytest
|
|
import scipy.misc
|
|
from scipy.fftpack import fftshift, ifftshift
|
|
|
|
|
|
|
|
@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 = cmath.complex_to_torch(obj)
|
|
|
|
# Construct wavefront from image
|
|
probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
|
|
return cmath.cmult(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(cmath.torch_to_complex(exit_waves_1)),norm='ortho'))
|
|
|
|
assert(np.allclose(np_result, cmath.torch_to_complex(propagators.far_field(exit_waves_1))))
|
|
|
|
|
|
|
|
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 = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho')))
|
|
|
|
assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
|
|
|
|
|
|
|
|
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(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
|
|
|
|
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))
|
|
|
|
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)
|
|
|
|
# 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(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_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.float64)
|
|
|
|
E0 = cmath.complex_to_torch(E0)
|
|
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))
|