diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index ae3d6bd..563d8f6 100644 --- a/tests/tools/test_interactions.py +++ b/tests/tools/test_interactions.py @@ -1,9 +1,10 @@ -from cdtools.tools import interactions import numpy as np -import torch as t -from numpy import fft from numpy.fft import fftshift, ifftshift +from numpy import fft import pytest +import torch as t + +from cdtools.tools import interactions # Have a random probe and a random object and test the two @@ -14,59 +15,60 @@ import pytest @pytest.fixture(scope='module') def random_probe(): - return np.random.rand(256,256) * np.exp(2j * np.pi * np.random.rand(256,256)) + return np.random.rand(256, 256) * np.exp(2j * np.pi * np.random.rand(256, 256)) + @pytest.fixture(scope='module') def random_obj(): - return np.random.rand(900,900) * np.exp(2j * np.pi * np.random.rand(900,900)) + return np.random.rand(900, 900) * np.exp(2j * np.pi * np.random.rand(900, 900)) + @pytest.fixture(scope='module') def single_pixel_probe(scope='module'): - probe = np.zeros((256,256), dtype=np.complex128) - probe[128,128] = 1 + probe = np.zeros((256, 256), dtype=np.complex128) + probe[128, 128] = 1 return probe def test_translations_to_pixel(): # 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)) + basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t() + translations = t.rand((10, 3)) output = interactions.translations_to_pixel(basis, translations) - assert t.allclose(output, -translations[:,:2].flip(1)) - + assert t.allclose(output, -translations[:, :2].flip(1)) + # Next, try a case with a single translation translation = t.rand((3)) output = interactions.translations_to_pixel(basis, translation) assert t.allclose(output, -translation[:2].flip(0)) - + # Then, try a case with no surface normal but with a real conversion - basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t() - translations = t.rand((10,3)) + basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t() + translations = t.rand((10, 3)) output = interactions.translations_to_pixel(basis, translations) basis_vectors_inv = t.pinverse(basis) - translations[:,2] = 0 # manually project off z component - assert t.allclose(output, t.mm(translations,basis_vectors_inv.t())) - + translations[:, 2] = 0 # manually project off z component + assert t.allclose(output, t.mm(translations, basis_vectors_inv.t())) + # Finally, try a case with a known surface normal (reflection) - basis = t.Tensor([[0,-1,0],[0,0,1]]).t() - surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)]) - translations = t.rand((10,3)) + basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t() + surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)]) + translations = t.rand((10, 3)) output = interactions.translations_to_pixel(basis, translations, surface_normal=surface_normal) - exp_translations = t.stack((-translations[:,1],translations[:,0]),dim=1) + exp_translations = t.stack((-translations[:, 1], translations[:, 0]), dim=1) assert t.allclose(output, exp_translations) - def test_pixel_to_translations(): # 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)) - translations[:,2] = 0 + basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t() + translations = t.rand((10, 3)) + translations[:, 2] = 0 output = interactions.translations_to_pixel(basis, translations) roundtrip = interactions.pixel_to_translations(basis, output) assert t.allclose(translations, roundtrip) - + # Next, try a case with a single translation translation = t.rand((3)) translation[2] = 0 @@ -74,65 +76,59 @@ def test_pixel_to_translations(): roundtrip = interactions.pixel_to_translations(basis, output) assert t.allclose(translation, roundtrip) - # Then, try a case with no surface normal but with a real conversion - basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t() - translations = t.rand((10,3)) - translations[:,2] = 0 # manually project off z component + basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t() + translations = t.rand((10, 3)) + translations[:, 2] = 0 # manually project off z component output = interactions.translations_to_pixel(basis, translations) roundtrip = interactions.pixel_to_translations(basis, output) assert t.allclose(translations, roundtrip) - + # Finally, try a case with a known surface normal (reflection) - basis = t.Tensor([[0,-1,0],[0,0,1]]).t() - surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)]) - translations = t.rand((10,3)) - translations[:,2] = 0 # manually project off z component + basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t() + surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)]) + translations = t.rand((10, 3)) + translations[:, 2] = 0 # manually project off z component output = interactions.translations_to_pixel(basis, translations, surface_normal=surface_normal) roundtrip = interactions.pixel_to_translations(basis, output, - surface_normal=surface_normal) + surface_normal=surface_normal) assert t.allclose(translations, roundtrip) - def test_project_translations_to_sample(): # 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)) + 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()) + 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)) + 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()) + 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 - translations = np.random.rand(10,2) * 500 - exit_waves_np = [random_probe * \ - random_obj[tr[0]:tr[0]+random_probe.shape[0], - tr[1]:tr[1]+random_probe.shape[1]] for + translations = np.random.rand(10, 2) * 500 + exit_waves_np = [random_probe * random_obj[tr[0]:tr[0] + random_probe.shape[0], + tr[1]:tr[1] + random_probe.shape[1]] for tr in np.round(translations).astype(int)] exit_waves_t = interactions.ptycho_2D_round(t.as_tensor(random_probe), t.as_tensor(random_obj), @@ -146,14 +142,13 @@ def test_ptycho_2D_round(random_probe, random_obj): assert np.allclose(exit_wave_t.numpy(), exit_waves_np[0]) - def test_ptycho_2D_linear(single_pixel_probe, random_obj): # For this one, I just want to check one translation, but # I need to check both formats - translations = np.array([[46.7,53.2]]) - translation = np.array([46.7,53.2]) - + translations = np.array([[46.7, 53.2]]) + translation = np.array([46.7, 53.2]) + exit_waves_probe = interactions.ptycho_2D_linear( t.as_tensor(single_pixel_probe), t.as_tensor(random_obj), @@ -167,14 +162,9 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj): shift_probe=True) # Check that the outputs match - assert t.allclose(exit_waves_probe[0],exit_wave_probe) + assert t.allclose(exit_waves_probe[0], exit_wave_probe) - - exit_waves_obj = interactions.ptycho_2D_linear( - t.as_tensor(single_pixel_probe), - t.as_tensor(random_obj), - t.tensor(translations), - shift_probe=False) + exit_waves_obj = interactions.ptycho_2D_linear(t.as_tensor(single_pixel_probe), t.as_tensor(random_obj), t.tensor(translations), shift_probe=False) exit_wave_obj = interactions.ptycho_2D_linear( t.as_tensor(single_pixel_probe), @@ -183,38 +173,37 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj): shift_probe=False) # Check that the outputs match - assert t.allclose(exit_waves_obj[0],exit_wave_obj) + assert t.allclose(exit_waves_obj[0], exit_wave_obj) # For the shifted probe, we should find 4 pixels with intensity exit_waves_probe = t.as_tensor(exit_waves_probe)[0] - probe_shift = np.array([[0.3*0.8,0.3*0.2], - [0.7*0.8,0.7*0.2]]) - obj_section = random_obj[128+46:128+48, - 128+53:128+55] - exit_section = exit_waves_probe[128:130,128:130] + probe_shift = np.array([[0.3 * 0.8, 0.3 * 0.2], + [0.7 * 0.8, 0.7 * 0.2]]) + obj_section = random_obj[128 + 46:128 + 48, + 128 + 53:128 + 55] + exit_section = exit_waves_probe[128:130, 128:130] assert np.allclose(probe_shift * obj_section, exit_section) - + # For the shifted obj, we should find one pixel with intensity exit_waves_obj = t.as_tensor(exit_waves_obj)[0] - obj_shift = np.array([[0.3*0.8,0.3*0.2], - [0.7*0.8,0.7*0.2]]) - obj_section = random_obj[128+46:128+48, - 128+53:128+55] - exit_pixel = exit_waves_obj[128,128] - assert np.isclose(np.sum(obj_shift * obj_section),exit_pixel) - + obj_shift = np.array([[0.3 * 0.8, 0.3 * 0.2], + [0.7 * 0.8, 0.7 * 0.2]]) + obj_section = random_obj[128 + 46:128 + 48, + 128 + 53:128 + 55] + exit_pixel = exit_waves_obj[128, 128] + assert np.isclose(np.sum(obj_shift * obj_section), exit_pixel) + # Test for a single translation def test_ptycho_2D_sinc(single_pixel_probe, random_obj): - - + # For this one, I just want to check one translation, but # I need to check both formats - translations = np.array([[46.7,53.2]]) - translation = np.array([46.7,53.2]) - + translations = np.array([[46.7, 53.2]]) + translation = np.array([46.7, 53.2]) + exit_waves_probe = interactions.ptycho_2D_sinc( t.as_tensor(single_pixel_probe), t.as_tensor(random_obj), @@ -228,32 +217,31 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj): shift_probe=True) # Check that the outputs match - assert t.allclose(exit_waves_probe[0],exit_wave_probe) - + assert t.allclose(exit_waves_probe[0], exit_wave_probe) # Now we explicitly define what the sinc interpolated array should # look like xs = np.arange(256) - 128 - Ys,Xs = np.meshgrid(xs,xs) + Ys, Xs = np.meshgrid(xs, xs) sinc_probe = np.sinc(Xs) * np.sinc(Ys) # Just check that the unshifted probe is correct assert np.allclose(single_pixel_probe, sinc_probe) - sinc_shifted_probe = np.sinc(Xs-0.7) * np.sinc(Ys-0.2) - obj_section = random_obj[46:46+256, - 53:53+256] + sinc_shifted_probe = np.sinc(Xs - 0.7) * np.sinc(Ys - 0.2) + obj_section = random_obj[46:46 + 256, + 53:53 + 256] exit_wave_np = sinc_shifted_probe * obj_section - + exit_wave_torch = exit_wave_probe.numpy() # The fidelity isn't great due to the FFT-based approach, so we need # a pretty relaxed condition - assert np.max(np.abs(exit_wave_np-exit_wave_torch)) < 0.005 + assert np.max(np.abs(exit_wave_np - exit_wave_torch)) < 0.005 def test_RPI_interaction(random_probe, random_obj): - random_obj1 = random_obj[:79,:68] * 0 + 1 + random_obj1 = random_obj[:79, :68] * 0 + 1 random_probe1 = random_probe * 0 + 1 t_random_obj1 = t.as_tensor(random_obj1) t_random_probe1 = t.as_tensor(random_probe1) @@ -261,36 +249,32 @@ def test_RPI_interaction(random_probe, random_obj): 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]//2 - 79//2: - -(random_probe1.shape[0]-79 - (random_probe1.shape[0]//2 - 79//2)), - (random_probe1.shape[1]-68)//2: - (random_probe1.shape[1]-68)//2 + 68] = obj1_fourier + obj1_ups[random_probe1.shape[0] // 2 - 79 // 2: + -(random_probe1.shape[0] - 79 - (random_probe1.shape[0] // 2 - 79 // 2)), + (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')) + norm='ortho')) - output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1] - / (random_obj1.shape[-2] * random_obj1.shape[-1])) + output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1] / (random_obj1.shape[-2] * random_obj1.shape[-1])) assert np.allclose(t_output1, output1) - - random_obj2 = np.stack([random_obj[:64,:89]]*3) - random_probe2 = random_probe[3:,5:] + + random_obj2 = np.stack([random_obj[:64, :89]] * 3) + random_probe2 = random_probe[3:, 5:] t_random_obj2 = t.as_tensor(random_obj2) t_random_probe2 = t.as_tensor(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 + 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')) + norm='ortho')) - output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1] - / (random_obj2.shape[-2] * random_obj2.shape[-1])) + output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1] / (random_obj2.shape[-2] * random_obj2.shape[-1])) - assert np.allclose(t_output2, output2) - diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index 633027d..b36b85d 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -1,79 +1,73 @@ -from cdtools.tools import losses import numpy as np import torch as t +from cdtools.tools import losses + # The idea here is to use a simple numpy calculation of the various # objective functions to check the torch implementations and make sure # that any optimizations in the future don't change the results + def test_amplitude_mse(): - # Make some fake data - data = np.random.rand(10,100,100) + data = np.random.rand(10, 100, 100) # And add some noise to it - sim = data + 0.1 * np.random.rand(10,100,100) + sim = data + 0.1 * np.random.rand(10, 100, 100) # and define a simple mask that needs to be broadcast - mask = (np.random.rand(100,100) > 0.1).astype(bool) + mask = (np.random.rand(100, 100) > 0.1).astype(bool) # First, test without a mask np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2) - #np_result /= data.size - torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim)) - assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + # np_result /= data.size + torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim)) + assert np.isclose(np_result, np.take(torch_result.numpy(), 0)) # Then, test with a mask np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2) - #np_result /= np.count_nonzero(mask * np.ones_like(data)) - torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim), - mask = t.from_numpy(mask)) - assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + # np_result /= np.count_nonzero(mask * np.ones_like(data)) + torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask)) + assert np.isclose(np_result, np.take(torch_result.numpy(), 0)) def test_intensity_mse(): # Make some fake data - data = np.random.rand(10,100,100) + data = np.random.rand(10, 100, 100) # And add some noise to it - sim = data + 0.1 * np.random.rand(10,100,100) + sim = data + 0.1 * np.random.rand(10, 100, 100) # and define a simple mask that needs to be broadcast - mask = (np.random.rand(100,100) > 0.1).astype(bool) - + mask = (np.random.rand(100, 100) > 0.1).astype(bool) # First, test without a mask np_result = np.sum((data - sim)**2) - np_result /= data.size - torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim)) - assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + np_result /= data.size + torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim)) + assert np.isclose(np_result, np.take(torch_result.numpy(), 0)) # Then, test with a mask np_result = np.sum(mask * (data - sim)**2) - np_result /= np.count_nonzero(mask * np.ones_like(data)) - torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim), - mask = t.from_numpy(mask)) - assert np.isclose(np_result, np.take(torch_result.numpy(),0)) - + np_result /= np.count_nonzero(mask * np.ones_like(data)) + torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask)) + assert np.isclose(np_result, np.take(torch_result.numpy(), 0)) + def test_poisson_nll(): # Make some fake data - data = np.random.rand(10,100,100) + data = np.random.rand(10, 100, 100) # And add some noise to it - sim = data + 0.1 * np.random.rand(10,100,100) + sim = data + 0.1 * np.random.rand(10, 100, 100) # and define a simple mask that needs to be broadcast - mask = (np.random.rand(100,100) > 0.1).astype(bool) - + mask = (np.random.rand(100, 100) > 0.1).astype(bool) # First, test without a mask np_result = np.sum(sim - data * np.log(sim)) - np_result /= data.size - torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim), eps=0) - assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + np_result /= data.size + torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim), eps=0) + assert np.isclose(np_result, np.take(torch_result.numpy(), 0)) # Then, test with a mask np_result = np.sum(mask * (sim - data * np.log(sim))) - np_result /= np.count_nonzero(mask * np.ones_like(data)) - torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim), - mask = t.from_numpy(mask), eps=0) - assert np.isclose(np_result, np.take(torch_result.numpy(),0)) - - - + np_result /= np.count_nonzero(mask * np.ones_like(data)) + torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim), + mask=t.from_numpy(mask), eps=0) + assert np.isclose(np_result, np.take(torch_result.numpy(), 0)) diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py index bfa06de..93ac72c 100644 --- a/tests/tools/test_measurements.py +++ b/tests/tools/test_measurements.py @@ -1,99 +1,90 @@ -from cdtools.tools import measurements import torch as t import numpy as np +from cdtools.tools import measurements + def test_intensity(): - wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10)) - epsilon=1e-6 + wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10)) + epsilon = 1e-6 np_result = np.abs(wavefields.numpy())**2 + epsilon - assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon), + assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon), t.as_tensor(np_result)) # Test single field case - assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon), + assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon), t.as_tensor(np_result[0])) - - det_slice = np.s_[3:,5:8] - assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon), - t.as_tensor(np_result[(np.s_[:],)+det_slice])) - + det_slice = np.s_[3:, 5:8] + assert t.allclose(measurements.intensity(wavefields, det_slice, epsilon=epsilon), + t.as_tensor(np_result[(np.s_[:],) + det_slice])) + # Test single field case - assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon), + assert t.allclose(measurements.intensity(wavefields[0], det_slice, epsilon=epsilon), t.as_tensor(np_result[0][det_slice])) - # With oversampling on - np_oversampling_result = (np_result[:,::2,::2] + \ - np_result[:,1::2,::2] + \ - np_result[:,::2,1::2] + \ - np_result[:,1::2,1::2]) / 4 + np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4 # With multiple fields - assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon, oversampling=2), + assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon, oversampling=2), t.as_tensor(np_oversampling_result,)) # With a single field - assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon, oversampling=2), + assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon, oversampling=2), t.as_tensor(np_oversampling_result[0],)) - + def test_incoherent_sum(): # With no explicit slice given - - wavefields = t.rand((5,4,10,10)) + 1j * t.rand((5,4,10,10)) - epsilon=1e-6 - np_result = np.sum(np.abs(wavefields.numpy())**2,axis=-3) + epsilon - assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon), + + wavefields = t.rand((5, 4, 10, 10)) + 1j * t.rand((5, 4, 10, 10)) + epsilon = 1e-6 + np_result = np.sum(np.abs(wavefields.numpy())**2, axis=-3) + epsilon + assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon), t.as_tensor(np_result)) # Test single field case - assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon), + assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon), t.as_tensor(np_result[0])) - # With a slice given - det_slice = np.s_[3:,5:8] - assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon), - t.as_tensor(np_result[(np.s_[:],)+det_slice])) + det_slice = np.s_[3:, 5:8] + assert t.allclose(measurements.incoherent_sum(wavefields, det_slice, epsilon=epsilon), + t.as_tensor(np_result[(np.s_[:],) + det_slice])) # Test single field case - assert t.allclose(measurements.incoherent_sum(wavefields[0,:],det_slice,epsilon=epsilon), + assert t.allclose(measurements.incoherent_sum(wavefields[0, :], det_slice, epsilon=epsilon), t.as_tensor(np_result[0][det_slice])) # With oversampling on - np_oversampling_result = (np_result[:,::2,::2] + \ - np_result[:,1::2,::2] + \ - np_result[:,::2,1::2] + \ - np_result[:,1::2,1::2]) / 4 + np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4 # With multiple fields - assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon, oversampling=2), + assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon, oversampling=2), t.as_tensor(np_oversampling_result,)) # With a single field - assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon, oversampling=2), + assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon, oversampling=2), t.as_tensor(np_oversampling_result[0],)) def test_quadratic_background(): # test with intensity - wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10)) - epsilon=1e-6 - background = t.rand((10,10)) + wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10)) + epsilon = 1e-6 + background = t.rand((10, 10)) np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon - det_slice = np.s_[3:,5:8] + det_slice = np.s_[3:, 5:8] - result = measurements.quadratic_background(wavefields,background[det_slice], + result = measurements.quadratic_background(wavefields, background[det_slice], detector_slice=det_slice, epsilon=epsilon, measurement=measurements.intensity) - assert t.allclose(result, t.tensor(np_result[(np.s_[:],)+det_slice])) - - + assert t.allclose(result, t.tensor(np_result[(np.s_[:],) + det_slice])) + # test with incoherent sum but no slice and no stack - wavefields = t.rand((4,10,10)) + 1j * t.rand((4,10,10)) - np_result = np.sum(np.abs(wavefields.numpy())**2,axis=0) + wavefields = t.rand((4, 10, 10)) + 1j * t.rand((4, 10, 10)) + np_result = np.sum(np.abs(wavefields.numpy())**2, axis=0) np_result += background.numpy()**2 result = measurements.quadratic_background(wavefields, background, epsilon=epsilon,