mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 13:32:40 +02:00
linting test_image_processing.py and test_initializers.py
This commit is contained in:
gnzng
@@ -1,18 +1,19 @@
|
||||
import numpy as np
|
||||
import torch as t
|
||||
from scipy import ndimage
|
||||
|
||||
from cdtools.tools import image_processing, interactions
|
||||
from scipy import ndimage
|
||||
|
||||
|
||||
def test_centroid():
|
||||
# Test single im
|
||||
im = t.rand((30,40))
|
||||
im = t.rand((30, 40))
|
||||
sp_centroid = ndimage.center_of_mass(im.numpy())
|
||||
centroid = image_processing.centroid(im)
|
||||
assert t.allclose(centroid, t.Tensor(sp_centroid))
|
||||
|
||||
|
||||
# Test stack o' ims
|
||||
ims = t.rand((5,30,40))
|
||||
ims = t.rand((5, 30, 40))
|
||||
sp_centroids = [ndimage.center_of_mass(im.numpy())
|
||||
for im in ims]
|
||||
centroids = image_processing.centroid(ims)
|
||||
@@ -21,33 +22,33 @@ def test_centroid():
|
||||
|
||||
def test_centroid_sq():
|
||||
# Test single im
|
||||
im = t.rand((30,40))
|
||||
im = t.rand((30, 40))
|
||||
sp_centroid = ndimage.center_of_mass(im.numpy()**2)
|
||||
centroid = image_processing.centroid_sq(im)
|
||||
assert t.allclose(centroid, t.Tensor(sp_centroid))
|
||||
|
||||
# Test complex with multiple ims
|
||||
ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40))
|
||||
ims = t.rand((5, 30, 40)) + 1j * t.rand((5, 30, 40))
|
||||
np_ims = ims.numpy()
|
||||
sp_centroids = [ndimage.center_of_mass(np.abs(im)**2)
|
||||
for im in np_ims]
|
||||
centroids = image_processing.centroid_sq(ims, comp=True)
|
||||
assert t.allclose(centroids, t.Tensor(np.array(sp_centroids)))
|
||||
|
||||
|
||||
|
||||
def test_sinc_subpixel_shift():
|
||||
|
||||
im = np.zeros((512,512), dtype=np.complex128)
|
||||
im[256,256] = 1
|
||||
im = np.zeros((512, 512), dtype=np.complex128)
|
||||
im[256, 256] = 1
|
||||
|
||||
# test it by creating a single pixel object and seeing that it is
|
||||
# shifted correctly
|
||||
xs = np.arange(512) - 256
|
||||
Ys,Xs = np.meshgrid(xs,xs)
|
||||
sinc_im = np.sinc(Xs-0.3) * np.sinc(Ys-0.6)
|
||||
Ys, Xs = np.meshgrid(xs, xs)
|
||||
sinc_im = np.sinc(Xs - 0.3) * np.sinc(Ys - 0.6)
|
||||
|
||||
torch_im = t.as_tensor(im)
|
||||
test_im = image_processing.sinc_subpixel_shift(torch_im,(0.3,0.6))
|
||||
test_im = image_processing.sinc_subpixel_shift(torch_im, (0.3, 0.6))
|
||||
|
||||
# The fidelity isn't great due to the FFT-based approach, so we need
|
||||
# a pretty relaxed condition
|
||||
@@ -57,84 +58,82 @@ def test_sinc_subpixel_shift():
|
||||
def test_find_pixel_shift():
|
||||
|
||||
# Test two real ims
|
||||
big_im = t.rand((30,70))
|
||||
im1 = big_im[3:,:-20]
|
||||
im2 = big_im[:-3,20:]
|
||||
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([-3,20]))
|
||||
big_im = t.rand((30, 70))
|
||||
im1 = big_im[3:, :-20]
|
||||
im2 = big_im[:-3, 20:]
|
||||
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([-3, 20]))
|
||||
|
||||
# Test a real and complex im
|
||||
big_im = t.rand((30,70))
|
||||
im1 = big_im[:-5,10:].to(dtype=t.complex64)
|
||||
im2 = big_im[5:,:-10]
|
||||
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10]))
|
||||
assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10]))
|
||||
|
||||
big_im = t.rand((30, 70))
|
||||
im1 = big_im[:-5, 10:].to(dtype=t.complex64)
|
||||
im2 = big_im[5:, :-10]
|
||||
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, -10]))
|
||||
assert t.all(image_processing.find_pixel_shift(im2, im1) == t.LongTensor([-5, 10]))
|
||||
|
||||
# Test two complex ims
|
||||
big_im = t.rand((45,45)) + 1j * t.rand((45,45))
|
||||
im1 = big_im[:-5,:-4]
|
||||
im2 = big_im[5:,4:]
|
||||
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4]))
|
||||
big_im = t.rand((45, 45)) + 1j * t.rand((45, 45))
|
||||
im1 = big_im[:-5, :-4]
|
||||
im2 = big_im[5:, 4:]
|
||||
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, 4]))
|
||||
|
||||
|
||||
def test_find_subpixel_shift():
|
||||
# We can do this by creating a test probe and a test object
|
||||
test_probe = t.rand((70,70)) + 1j * t.rand((70,70))
|
||||
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
|
||||
test_probe = t.rand((70, 70)) + 1j * t.rand((70, 70))
|
||||
test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
|
||||
|
||||
shift = t.tensor((0.8, 0.75))
|
||||
|
||||
shift = t.tensor((0.8,0.75))
|
||||
|
||||
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, multiple_modes=False)
|
||||
|
||||
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50)
|
||||
|
||||
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0, 0), resolution=50)
|
||||
# tolerance of 0.03 on this measurement
|
||||
assert t.all(t.abs(shift - retrieved_shift) < 0.03)
|
||||
|
||||
|
||||
|
||||
def test_find_shift():
|
||||
|
||||
# We can do this by creating a test probe and a test object
|
||||
test_probe = t.rand((200,200)) + 1j * t.rand((200,200))
|
||||
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
|
||||
test_probe = t.rand((200, 200)) + 1j * t.rand((200, 200))
|
||||
test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
|
||||
|
||||
shift = t.tensor((0.8, 0.75))
|
||||
|
||||
shift = t.tensor((0.8,0.75))
|
||||
|
||||
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift,
|
||||
multiple_modes=False)[:-40,:-6]
|
||||
multiple_modes=False)[:-40, :-6]
|
||||
|
||||
retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50)
|
||||
retrieved_shift = image_processing.find_shift(im, test_probe[40:, 6:], resolution=50)
|
||||
# tolerance of 0.03 on this measurement
|
||||
assert t.all(t.abs(shift + t.Tensor((40,6)) - retrieved_shift) < 0.03)
|
||||
assert t.all(t.abs(shift + t.Tensor((40, 6)) - retrieved_shift) < 0.03)
|
||||
|
||||
|
||||
|
||||
def test_convolve_1d():
|
||||
test_image = np.random.rand(400,300)
|
||||
#test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
|
||||
xs = np.linspace(-100,100,300)
|
||||
kernel = 1/(1+xs**2)
|
||||
test_image = np.random.rand(400, 300)
|
||||
# test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
|
||||
xs = np.linspace(-100, 100, 300)
|
||||
kernel = 1 / (1 + xs**2)
|
||||
|
||||
# First, we test with everything real, dim=1
|
||||
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
|
||||
t.as_tensor(kernel),dim=1)
|
||||
t.as_tensor(kernel), dim=1)
|
||||
|
||||
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
|
||||
assert np.allclose(convolved.numpy(),np_result)
|
||||
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
|
||||
assert np.allclose(convolved.numpy(), np_result)
|
||||
|
||||
|
||||
xs = np.linspace(-100,100,400)
|
||||
kernel = 1/(1+xs**2)
|
||||
xs = np.linspace(-100, 100, 400)
|
||||
kernel = 1 / (1 + xs**2)
|
||||
|
||||
# Then with dim=0, and a non-fftshifted kernel
|
||||
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
|
||||
t.as_tensor(np.fft.ifftshift(kernel)),
|
||||
fftshift_kernel=False)
|
||||
|
||||
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0))
|
||||
assert np.allclose(convolved.numpy(),np_result)
|
||||
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0))
|
||||
assert np.allclose(convolved.numpy(), np_result)
|
||||
|
||||
# And finally with complex input
|
||||
convolved = image_processing.convolve_1d(t.as_tensor(test_image,dtype=t.complex64),
|
||||
t.as_tensor(kernel,dtype=t.complex64)).numpy()
|
||||
|
||||
np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)
|
||||
assert np.allclose(convolved,np_result)
|
||||
convolved = image_processing.convolve_1d(t.as_tensor(test_image, dtype=t.complex64),
|
||||
t.as_tensor(kernel, dtype=t.complex64)).numpy()
|
||||
|
||||
np_result = np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0)
|
||||
assert np.allclose(convolved, np_result)
|
||||
|
||||
@@ -1,26 +1,27 @@
|
||||
from cdtools.tools import initializers
|
||||
from cdtools.datasets import Ptycho2DDataset
|
||||
import numpy as np
|
||||
import torch as t
|
||||
|
||||
from cdtools.tools import initializers
|
||||
from cdtools.datasets import Ptycho2DDataset
|
||||
|
||||
|
||||
def test_exit_wave_geometry():
|
||||
# First test a simple case where nothing need change
|
||||
basis = t.Tensor([[0,-30e-6,0],
|
||||
[-20e-6,0,0]]).transpose(0,1)
|
||||
shape = t.Size([73,56])
|
||||
basis = t.Tensor([[0, -30e-6, 0],
|
||||
[-20e-6, 0, 0]]).transpose(0, 1)
|
||||
shape = t.Size([73, 56])
|
||||
wavelength = 1e-9
|
||||
distance = 1.
|
||||
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
|
||||
|
||||
assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07]))
|
||||
assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07]))
|
||||
|
||||
|
||||
assert t.allclose(rs_basis[0, 1], t.Tensor([-8.928571428571428e-07]))
|
||||
assert t.allclose(rs_basis[1, 0], t.Tensor([-4.5662100456621004e-07]))
|
||||
|
||||
|
||||
def test_calc_object_setup():
|
||||
# First just try a simple case
|
||||
probe_shape = t.Size([120,57])
|
||||
translations = t.rand((30,2)) * 300
|
||||
probe_shape = t.Size([120, 57])
|
||||
translations = t.rand((30, 2)) * 300
|
||||
t_max = t.max(translations, dim=0)[0]
|
||||
t_min = t.min(translations, dim=0)[0]
|
||||
obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations)
|
||||
@@ -28,65 +29,58 @@ def test_calc_object_setup():
|
||||
assert t.allclose(min_translation, t_min)
|
||||
|
||||
assert obj_shape == t.Size(exp_shape)
|
||||
|
||||
|
||||
# Then add some padding
|
||||
padding = 5
|
||||
obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations, padding=padding)
|
||||
assert t.allclose(min_translation, t_min - padding)
|
||||
assert obj_shape == t.Size(exp_shape + 2 * padding)
|
||||
|
||||
|
||||
|
||||
|
||||
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)
|
||||
|
||||
center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
|
||||
y, x = np.mgrid[:shape[0], :shape[1]]
|
||||
np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
|
||||
-0.5*((y-center[0])/sigma[0])**2)
|
||||
np_result = 10 * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
|
||||
init_result = initializers.gaussian(shape, sigma, amplitude=10).numpy()
|
||||
assert np.allclose(init_result, np_result)
|
||||
|
||||
# Generate gaussian as a numpy array (rectangular array)
|
||||
shape = [10, 5]
|
||||
sigma = [2.5, 3]
|
||||
center = ((shape[0]-1)/2, (shape[1]-1)/2)
|
||||
center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
|
||||
y, x = np.mgrid[:shape[0], :shape[1]]
|
||||
np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2
|
||||
-0.5*((y-center[0])/sigma[0])**2)
|
||||
np_result = np.exp(-0.5 * ((x - center[1]) / sigma[1])**2
|
||||
- 0.5 * ((y - center[0]) / sigma[0])**2)
|
||||
init_result = initializers.gaussian(shape, sigma).numpy()
|
||||
assert np.allclose(init_result, np_result)
|
||||
|
||||
|
||||
# Generate gaussian with curvature
|
||||
shape = [20, 30]
|
||||
sigma = [2.5, 5]
|
||||
curvature = [1,0.6]
|
||||
center = ((shape[0]-1)/2 + 3, (shape[1]-1)/2 - 1.4)
|
||||
curvature = [1, 0.6]
|
||||
center = ((shape[0] - 1) / 2 + 3, (shape[1] - 1) / 2 - 1.4)
|
||||
y, x = np.mgrid[:shape[0], :shape[1]]
|
||||
np_result = (10+0j)*np.exp(-0.5*((x-center[1])/sigma[1])**2
|
||||
-0.5*((y-center[0])/sigma[0])**2)
|
||||
np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2
|
||||
+0.5j*curvature[0]*(y-center[0])**2)
|
||||
init_result = initializers.gaussian(shape, sigma, center=center,
|
||||
curvature=curvature, amplitude=10).numpy()
|
||||
np_result = (10 + 0j) * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
|
||||
np_result *= np.exp(0.5j * curvature[1] * (x - center[1])**2 + 0.5j * curvature[0] * (y - center[0])**2)
|
||||
init_result = initializers.gaussian(shape, sigma, center=center, curvature=curvature, amplitude=10).numpy()
|
||||
assert np.allclose(init_result, np_result)
|
||||
|
||||
|
||||
|
||||
def test_gaussian_probe(ptycho_cxi_1):
|
||||
|
||||
dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
|
||||
|
||||
det_basis = t.Tensor(dataset.detector_geometry['basis'])
|
||||
det_shape = t.Size(dataset.patterns.shape[-2:])
|
||||
wavelength = dataset.wavelength
|
||||
distance = dataset.detector_geometry['distance']
|
||||
|
||||
basis = initializers.exit_wave_geometry(det_basis,
|
||||
det_shape,
|
||||
wavelength,
|
||||
distance)
|
||||
det_shape,
|
||||
wavelength,
|
||||
distance)
|
||||
# Basis is around 60nm in the i(y) direction, 85nm in the j(x) direction
|
||||
# Full window is therefore about 15 um in i(y) and 20 um in the j(x) dir
|
||||
|
||||
@@ -94,15 +88,13 @@ def test_gaussian_probe(ptycho_cxi_1):
|
||||
sigma = 5e-7
|
||||
|
||||
# Build a stage explicitly with numpy to compare against
|
||||
x = (np.arange(256) - 127.5) * (-basis[0,1]).numpy()
|
||||
y = (np.arange(256) - 127.5) * (-basis[1,0]).numpy()
|
||||
Xs,Ys = np.meshgrid(x,y)
|
||||
Rs = np.sqrt(Xs**2+Ys**2)
|
||||
|
||||
|
||||
x = (np.arange(256) - 127.5) * (-basis[0, 1]).numpy()
|
||||
y = (np.arange(256) - 127.5) * (-basis[1, 0]).numpy()
|
||||
Xs, Ys = np.meshgrid(x, y)
|
||||
Rs = np.sqrt(Xs**2 + Ys**2)
|
||||
|
||||
# Now we first test the non-propagated probe
|
||||
np_probe = np.exp(-1/(2*sigma**2) * Rs**2)
|
||||
np_probe = np.exp(- 1 / (2 * sigma**2) * Rs**2)
|
||||
|
||||
normalization = 0
|
||||
for params, im in dataset:
|
||||
@@ -110,27 +102,26 @@ def test_gaussian_probe(ptycho_cxi_1):
|
||||
normalization /= len(dataset)
|
||||
|
||||
normalization_1 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
|
||||
|
||||
|
||||
probe = initializers.gaussian_probe(
|
||||
dataset, basis, det_shape, sigma).numpy()
|
||||
|
||||
assert np.allclose(probe, normalization_1*np_probe)
|
||||
|
||||
assert np.allclose(probe, normalization_1 * np_probe)
|
||||
|
||||
# And then a propagated probe
|
||||
z = 1e-4 #nm
|
||||
z = 1e-4 # nm
|
||||
k = 2 * np.pi / wavelength
|
||||
w0 = np.sqrt(2)*sigma
|
||||
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)
|
||||
np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
|
||||
Rz = z * (1 + (zr / z)**2)
|
||||
np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
|
||||
|
||||
normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
|
||||
|
||||
normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
|
||||
|
||||
probe = initializers.gaussian_probe(dataset, basis, det_shape, sigma,
|
||||
propagation_distance=z).numpy()
|
||||
|
||||
assert np.allclose(probe, normalization_2*np_probe)
|
||||
assert np.allclose(probe, normalization_2 * np_probe)
|
||||
|
||||
|
||||
def test_SHARP_style_probe(ptycho_cxi_1):
|
||||
@@ -148,11 +139,13 @@ def test_SHARP_style_probe(ptycho_cxi_1):
|
||||
wavelength,
|
||||
distance)
|
||||
|
||||
assert basis.shape == t.Size([3, 2])
|
||||
|
||||
probe = initializers.SHARP_style_probe(dataset)
|
||||
assert probe.shape == t.Size([256,256])
|
||||
assert probe.shape == t.Size([256, 256])
|
||||
|
||||
probe = initializers.SHARP_style_probe(dataset, propagation_distance=20e-6)
|
||||
assert probe.shape == t.Size([256,256])
|
||||
assert probe.shape == t.Size([256, 256])
|
||||
|
||||
|
||||
def test_RPI_spectral_init():
|
||||
@@ -160,28 +153,27 @@ def test_RPI_spectral_init():
|
||||
# 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]
|
||||
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.as_tensor(np.random.rand(*pattern.shape),dtype=t.float32) * 0.05
|
||||
background = t.as_tensor(np.random.rand(*pattern.shape), dtype=t.float32) * 0.05
|
||||
|
||||
probe = t.as_tensor(probe)
|
||||
pattern = t.as_tensor(pattern)
|
||||
|
||||
|
||||
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
|
||||
assert list(obj.shape) == [1]+obj_shape
|
||||
assert list(obj.shape) == [1] + obj_shape
|
||||
|
||||
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
|
||||
n_modes=2, mask=mask)
|
||||
assert list(obj.shape) == [2]+obj_shape
|
||||
assert list(obj.shape) == [2] + obj_shape
|
||||
|
||||
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
|
||||
n_modes=2, background=background)
|
||||
assert list(obj.shape) == [2]+obj_shape
|
||||
assert list(obj.shape) == [2] + obj_shape
|
||||
|
||||
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
|
||||
n_modes=2, mask=mask,
|
||||
background=background)
|
||||
assert list(obj.shape) == [2]+obj_shape
|
||||
|
||||
assert list(obj.shape) == [2] + obj_shape
|
||||
|
||||
Reference in New Issue
Block a user