mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
linting test_interactions.py, test_losses.py, and test_measurements.py
This commit is contained in:
gnzng
@@ -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)
|
||||
|
||||
|
||||
+32
-38
@@ -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))
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user