Maybe last time it didn't actually commit?

tests/tools/test_analysis.py

@@ -22,8 +22,8 @@ def test_orthogonalize_probes():
     probe_Ys, probe_Xs = np.meshgrid(probe_ys, probe_xs)
     probe_Rs = np.sqrt(probe_Xs**2 + probe_Ys**2)
 
-    probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2)),
-                       3*np.exp(-probe_Rs**2 / (2 * 12**2)),
+    probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2 + 1j)),
+                       3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)),
                        1*np.exp(-probe_Rs**2 / (2 * 15**2))]).astype(np.complex64)
 
     # test that it works on numpy arrays
@@ -31,13 +31,15 @@ def test_orthogonalize_probes():
 
     # test that it also works on torch tensors
     ortho_probes_t = cmath.torch_to_complex(analysis.orthogonalize_probes(cmath.complex_to_torch(probes)))
-        
+
+    # This tests for orthogonality
     for p1,p2 in combinations(ortho_probes,2):
         assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
 
     for p1,p2 in combinations(ortho_probes_t,2):
         assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
 
+        
     probe_intensity = np.sum(np.abs(probes)**2,axis=0)
     ortho_probe_intensity = np.sum(np.abs(ortho_probes)**2,axis=0)
     ortho_probe_t_intensity = np.sum(np.abs(ortho_probes_t)**2,axis=0)
@@ -45,8 +47,35 @@ def test_orthogonalize_probes():
     assert np.allclose(probe_intensity,ortho_probe_intensity)
     assert np.allclose(probe_intensity,ortho_probe_t_intensity)
 
+    # Check that it returns normalized probes if we ask
+    ortho_probes = analysis.orthogonalize_probes(probes, normalize=True)
+    assert np.allclose([1,1,1],np.sum(np.abs(ortho_probes)**2,axis=(1,2)))
 
+    # And now we check that the A matrices actually work
+    ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True,
+                                                    normalize=False)
+    assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes)
 
+    # And now we check that the A matrices actually work
+    ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True,
+                                                    normalize=False)
+    assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes)
+
+    # The big problem here is that we haven't tested if it returns the
+    # probes or their complex conjugate.. we can test that by sending in
+    # one probe and checking that we get the right thing out
+    
+    probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 20**2 + 100j))])
+    ortho_probes = analysis.orthogonalize_probes(probes)
+    # We need to correct for the arbitrary phase offset to be able to
+    # compare them
+    correction = ortho_probes[0]/probes[0]
+    correction = correction / np.abs(correction)
+    ortho_probes *= correction
+
+    assert np.allclose(ortho_probes,probes)
+
+        
 def test_standardize():
 
     # Start by making a probe and object that should meet the standardization

tests/tools/test_initializers.py

@@ -192,6 +192,32 @@ def test_SHARP_style_probe(ptycho_cxi_1):
 
 
 def test_RPI_spectral_init():
-    # Figure out a good way to test this
-    assert 0
-    RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, background=None)
+    # I think we can only really meaningfully test that it doesn't throw errors,
+    # since the original implementation is in numpy and there aren't any clear
+    # cases that can be calculated analytically.
+
+    pattern = np.random.rand(230,253).astype(np.float32)
+    probe = np.random.rand(230,253).astype(np.complex64)
+    obj_shape = [37,53]
+    mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04)
+    background = t.Tensor(np.random.rand(*pattern.shape) .astype(np.float32)* 0.05)
+
+    probe = cmath.complex_to_torch(probe)
+    pattern = t.Tensor(pattern)
+    
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
+    assert list(obj.shape) == [1]+obj_shape+[2]
+
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
+                                         n_modes=2, mask=mask)
+    assert list(obj.shape) == [2]+obj_shape+[2]
+
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
+                                         n_modes=2, background=background)
+    assert list(obj.shape) == [2]+obj_shape+[2]
+
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
+                                         n_modes=2, mask=mask,
+                                         background=background)
+    assert list(obj.shape) == [2]+obj_shape+[2]
+

tests/tools/test_interactions.py

@@ -4,6 +4,8 @@ from CDTools.tools import cmath
 from CDTools.tools import interactions
 import numpy as np
 import torch as t
+from numpy import fft
+from scipy.fftpack import fftshift, ifftshift
 import pytest
 
 
@@ -98,8 +100,35 @@ def test_pixel_to_translations():
 
 
 def test_project_translations_to_sample():
-    # Needs to be tested
-    assert 0
+    # First, try the case where everything is ones and simple
+    basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
+    translations = t.rand((10,3))
+    pixels, props = interactions.project_translations_to_sample(basis, translations)
+
+    assert np.allclose(pixels[:,0].numpy(),-translations[:,1])
+    assert np.allclose(pixels[:,1].numpy(),-translations[:,0])
+    assert np.allclose(props.numpy(),-translations[:,2:].numpy())
+
+    # Next, a simple tilt along one axis. This is a 45 degree rotation
+    # around the positive y-axis
+    # Thus, y-axis translations are unaffected, but x-axis translations
+    # induce a motion of 1/sqrt(2) in the j- pixel space, as well as
+    # creating a propagation (negative propagation for positive x)
+    basis = t.Tensor([[0,-1e-3,0],[-np.sqrt(2)*1e-3,0,np.sqrt(2)*1e-3]]).t()
+    translations = t.rand((10,3))
+    pixels, props = interactions.project_translations_to_sample(basis, translations)
+
+    print(props.numpy())
+    print(-translations[:,2:].numpy() - translations[:,:1].numpy())
+    assert np.allclose(pixels[:,0].numpy(),-translations[:,1]*1e3)
+    assert np.allclose(pixels[:,1].numpy(),-translations[:,0]*1e3/np.sqrt(2))
+    assert np.allclose(props.numpy(),-translations[:,2:].numpy() - translations[:,:1].numpy())
+
+    # Finally, we check a non-orthogonal case
+
+
+    
+    
 
 def test_ptycho_2D_round(random_probe, random_obj):
     # Test a stack of images
@@ -227,22 +256,39 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
 
 def test_RPI_interaction(random_probe, random_obj):
 
-    random_obj1 = cmath.complex_to_torch(random_obj[:79,:68])
-    random_probe1 = cmath.complex_to_torch(random_probe)
-    output = interactions.RPI_interaction(random_probe1, random_obj1)
 
+    random_obj1 = random_obj[:79,:68]
+    random_probe1 = random_probe
+    t_random_obj1 = cmath.complex_to_torch(random_obj1)
+    t_random_probe1 = cmath.complex_to_torch(random_probe1)
+    t_output1 = interactions.RPI_interaction(t_random_probe1, t_random_obj1)
 
-    random_obj1 = cmath.complex_to_torch(random_obj[:256,:256])
-    random_probe1 = cmath.complex_to_torch(random_probe)
-    output = interactions.RPI_interaction(random_probe1, random_obj1)
+    obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho'))
+    obj1_ups = np.zeros(random_probe1.shape[:2]).astype(np.complex128)
+    obj1_ups[(random_probe1.shape[0]-79)//2:
+             (random_probe1.shape[0]-79)//2 + 79,
+             (random_probe1.shape[1]-68)//2:
+             (random_probe1.shape[1]-68)//2 + 68] = obj1_fourier
+    output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
+                                                norm='ortho'))
+    
+    assert np.allclose(cmath.torch_to_complex(t_output1), output1)
 
     
-    random_obj1 = cmath.complex_to_torch(random_obj[:42,:103])
-    random_probe1 = cmath.complex_to_torch(random_probe)
-    output = interactions.RPI_interaction(random_probe1, random_obj1)
+    random_obj2 = np.stack([random_obj[:64,:89]]*3)
+    random_probe2 = random_probe[3:,5:]
+    t_random_obj2 = cmath.complex_to_torch(random_obj2)
+    t_random_probe2 = cmath.complex_to_torch(random_probe2)
+    t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2)
+
+    obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho'))
+    obj2_ups = np.zeros((3,)+random_probe2.shape[:2]).astype(np.complex128)
+    obj2_ups[:,(random_probe2.shape[0]-64)//2:
+             (random_probe2.shape[0]-64)//2 + 64,
+             (random_probe2.shape[1]-89)//2:
+             (random_probe2.shape[1]-89)//2 + 89] = obj2_fourier
+    output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups),
+                                                norm='ortho'))
+
 
-    # Need to actually test against a numpy implementation
-    print(random_probe.shape)
-    print(random_obj.shape)
-    assert 0