From 88df4c079484f401bcb5812d0eabef2b4d952f8a Mon Sep 17 00:00:00 2001 From: gnzng Date: Mon, 7 Jul 2025 14:40:17 -0700 Subject: [PATCH] linting test_image_processing.py and test_initializers.py --- tests/tools/test_image_processing.py | 115 ++++++++++++------------- tests/tools/test_initializers.py | 124 +++++++++++++-------------- 2 files changed, 115 insertions(+), 124 deletions(-) diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index 83c34d8..a6385cd 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -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) diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index 53c284a..5430271 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -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