Checked and fixed the tests, got them all running

This commit is contained in:
Abe Levitan
2024-02-12 14:55:26 -05:00
parent 1502d567bf
commit 5925ee74ab
7 changed files with 24 additions and 48 deletions
Binary file not shown.
+1 -3
View File
@@ -25,9 +25,7 @@ model = cdtools.models.Bragg2DPtycho.from_dataset(
#model.to(device='cuda')
#dataset.get_as(device='cuda')
model.inspect(dataset)
plt.show()
exit()
model.translation_offsets.requires_grad = False
for loss in model.Adam_optimize(100, dataset):
model.inspect(dataset)
@@ -21,7 +21,7 @@ __all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
'RPI_spectral_init',
'generate_subdominant_modes']
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, padding=0, oversampling=1):
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, oversampling=1):
"""Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry
It takes in the parameters for a given detector - the basis defining
-3
View File
@@ -138,6 +138,3 @@ def test_convolve_1d():
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)
def test_center():
raise Exception('Not yet tested')
+13 -29
View File
@@ -15,23 +15,6 @@ def test_exit_wave_geometry():
assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07]))
assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07]))
# Then test it's padding function
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength,
distance, padding=2)
exp_shape = t.Size([77,60])
assert full_shape == exp_shape
assert t.ones(full_shape)[det_slice].shape == shape
# Finally test it off-center, without expanding
center = t.Tensor([20,42])
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, center=center)
exp_shape = t.Size([105,84])
assert full_shape == exp_shape
assert t.ones(full_shape)[det_slice].shape == shape
def test_calc_object_setup():
@@ -100,10 +83,10 @@ def test_gaussian_probe(ptycho_cxi_1):
wavelength = dataset.wavelength
distance = dataset.detector_geometry['distance']
basis, shape, s = initializers.exit_wave_geometry(det_basis,
det_shape,
wavelength,
distance)
basis = initializers.exit_wave_geometry(det_basis,
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
@@ -128,7 +111,8 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization_1 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
probe = initializers.gaussian_probe(dataset, basis, shape, sigma).numpy()
probe = initializers.gaussian_probe(
dataset, basis, det_shape, sigma).numpy()
assert np.allclose(probe, normalization_1*np_probe)
@@ -143,7 +127,7 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
probe = initializers.gaussian_probe(dataset, basis, shape, sigma,
probe = initializers.gaussian_probe(dataset, basis, det_shape, sigma,
propagation_distance=z).numpy()
assert np.allclose(probe, normalization_2*np_probe)
@@ -159,15 +143,15 @@ def test_SHARP_style_probe(ptycho_cxi_1):
wavelength = dataset.wavelength
distance = dataset.detector_geometry['distance']
basis, shape, det_slice = initializers.exit_wave_geometry(det_basis,
det_shape,
wavelength,
distance)
basis = initializers.exit_wave_geometry(det_basis,
det_shape,
wavelength,
distance)
probe = initializers.SHARP_style_probe(dataset, shape, det_slice)
probe = initializers.SHARP_style_probe(dataset)
assert probe.shape == t.Size([256,256])
probe = initializers.SHARP_style_probe(dataset, shape, det_slice, propagation_distance=20e-6)
probe = initializers.SHARP_style_probe(dataset, propagation_distance=20e-6)
assert probe.shape == t.Size([256,256])
+4 -4
View File
@@ -8,7 +8,7 @@ import matplotlib.pyplot as plt
def test_plot_amplitude(show_plot):
# Test with tensor
im = t.as_tensor(scipy.datasets.ascent(),dtype=t.complex128)
plotting.plot_amplitude(im, basis = np.array([[1,1], [1,1], [0,0]]), title = 'Test Amplitude')
plotting.plot_amplitude(im, basis = np.array([[0,-1], [-1,0], [0,0]]), title = 'Test Amplitude')
if show_plot:
plt.show()
@@ -28,15 +28,15 @@ def test_plot_phase(show_plot):
# Test with numpy array
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy()
plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[1,1], [1,1], [0,0]]))
plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[0,-1], [-1,0], [0,0]]))
if show_plot:
plt.show()
def test_plot_colorize(show_plot):
def test_plot_colorized(show_plot):
# Test with tensor
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64)
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[1,1], [1,1], [0,0]]))
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[0,-1], [-1,0], [0,0]]))
if show_plot:
plt.show()
+5 -8
View File
@@ -48,12 +48,10 @@ def test_generate_high_NA_k_intensity_map():
basis = t.Tensor([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(0,1)
shape = t.Size([478,573])
#shape = t.Size([3,5])
wavelength = 1e-9
distance = 1#6e-3
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance)
distance = 1
rs_basis = \
initializers.exit_wave_geometry(basis, shape, wavelength, distance)
k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
rs_basis, basis, shape, distance, wavelength,
@@ -82,9 +80,8 @@ def test_generate_high_NA_k_intensity_map():
#print(rs_basis)
#print(rs_basis_tilted)
distance = 0.01#6e-3
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance)
rs_basis = \
initializers.exit_wave_geometry(basis, shape, wavelength, distance)
rs_basis_tilted = rs_basis.clone()
rs_basis_tilted[2,1] = rs_basis_tilted[0,1]