add projectors, propagators, and initializers

tests/tools/test_initializers.py

@@ -0,0 +1,24 @@
+from __future__ import division, print_function, absolute_import
+
+from CDTools.tools import initializers
+import numpy as np
+import torch as t
+
+
+
+def test_gaussian():
+    # Generate gaussian as a numpy array (square array)
+    shape = [10, 10]
+    sigma = [2.5, 2.5]
+    center = ((shape[0]-1)/2, (shape[1]-1)/2)
+    y, x = np.mgrid[:shape[0], :shape[1]]
+    np_result = 10*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**2)
+    assert(np.allclose(initializers.gaussian([10, 10], 10, [2.5, 2.5]), np_result))
+
+    # Generate gaussian as a numpy array (rectangular array)
+    shape = [10, 5]
+    sigma = [2.5, 2.5]
+    center = ((shape[0]-1)/2, (shape[1]-1)/2)
+    y, x = np.mgrid[:shape[0], :shape[1]]
+    np_result = 10*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**2)
+    assert(np.allclose(initializers.gaussian([10, 5], 10, [2.5, 2.5]), np_result))

tests/tools/test_projectors.py

@@ -0,0 +1,21 @@
+from __future__ import division, print_function, absolute_import
+
+from CDTools.tools import cmath
+from CDTools.tools import projectors
+import numpy as np
+import torch as t
+from scipy.fftpack import fftshift, ifftshift
+
+def test_modulus():
+    # Create a complex array with modulus 12 and phase pi/4
+    np_result = 6**.5*(np.ones((10,10))+1j*np.ones((10,10)))
+
+    assert(np.allclose(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)))), np_result))
+
+
+def test_support():
+    # Test masking
+    support = t.zeros((10,10))
+    np_result = np.zeros((10,10))
+
+    assert(np.allclose(cmath.torch_to_complex(projectors.support(t.ones((10,10,2)), support)), np_result))

tests/tools/test_propagators.py

@@ -0,0 +1,48 @@
+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
+    object = scipy.misc.ascent()[0:64,0:64].astype(np.complex128)
+    arr = np.random.random_sample((64,64))
+    object *= (arr+(1-arr**2)**.5*1j)
+
+    # Construct wavefront from image
+    probe = initializers.gaussian([64, 64], 1e3, [5, 5])*(1+1j)
+    return cmath.complex_to_torch(probe*object)
+
+
+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(cmath.torch_to_complex(exit_waves_1)))
+
+    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
+    np_result = exit_waves_1
+    # Far field result for exit waves calculated with numpy
+    far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1))))
+
+    assert(np.allclose(np_result, propagators.inverse_far_field(far_field_np_result)))
+
+
+
+def test_near_field(exit_waves_1):
+    pass
+
+
+
+def test_get_exit_waves(exit_waves_1):
+    pass