def test_ssb(dpix, backend, n_threads):
lt_ctx = lt.Context(InlineJobExecutor(debug=True, inline_threads=n_threads))
try:
if backend == 'cupy':
set_use_cuda(0)
dtype = np.float64
scaling = 4
shape = (29, 30, 189 // scaling, 197 // scaling)
# The acceleration voltage U in keV
U = 300
lamb = wavelength(U)
# STEM semiconvergence angle in radians
semiconv = 25e-3
# Diameter of the primary beam in the diffraction pattern in pixels
semiconv_pix = 78.6649 / scaling
cy = 93 // scaling
cx = 97 // scaling
input_data = (
np.random.uniform(0, 1, np.prod(shape))
* np.linspace(1.0, 1000.0, num=np.prod(shape))
)
input_data = input_data.astype(np.float64).reshape(shape)
udf = SSB_UDF(lamb=lamb, dpix=dpix, semiconv=semiconv, semiconv_pix=semiconv_pix,
dtype=dtype, cy=cy, cx=cx, method='subpix')
dataset = MemoryDataSet(
data=input_data, tileshape=(20, shape[2], shape[3]), num_partitions=2, sig_dims=2,
)
result = lt_ctx.run_udf(udf=udf, dataset=dataset)
result_f, reference_masks = reference_ssb(input_data, U=U, dpix=dpix, semiconv=semiconv,
semiconv_pix=semiconv_pix, cy=cy, cx=cx)
task_data = udf.get_task_data()
udf_masks = task_data['masks'].computed_masks
half_y = shape[0] // 2 + 1
# Use symmetry and reshape like generate_masks()
reference_masks = reference_masks[:half_y].reshape((half_y*shape[1], shape[2], shape[3]))
print(np.max(np.abs(udf_masks.todense() - reference_masks)))
print(np.max(np.abs(result['fourier'].data - result_f)))
assert np.allclose(result['fourier'].data, result_f)
backwards = result['amplitude'].data**2 * np.exp(1j*result['phase'].data)
assert np.allclose(result['fourier'].data, np.fft.fft2(backwards))
finally:
if backend == 'cupy':
set_use_cpu(0)
def test_masks():
scaling = 4
dtype = np.float64
shape = (29, 30, 189 // scaling, 197 // scaling)
# ? shape = np.random.uniform(1, 300, (4,1,))
# The acceleration voltage U in keV
U = 300
lambda_e = wavelength(U)
# STEM pixel size in m, here 50 STEM pixels on 0.5654 nm
dpix = 0.5654/50*1e-9
# STEM semiconvergence angle in radians
semiconv = 25e-3
# Diameter of the primary beam in the diffraction pattern in pixels
semiconv_pix = 78.6649 / scaling
cy = 93 // scaling
cx = 97 // scaling
input_data = np.random.uniform(0, 1, shape)
_, reference_masks = reference_ssb(
input_data, U=U, dpix=dpix, semiconv=semiconv,
semiconv_pix=semiconv_pix,
cy=cy, cx=cx
)
# print(np.max(np.abs(np.abs(result['fourier']) - np.abs(result_f))))
half_y = shape[0] // 2 + 1
# Use symmetry and reshape like generate_masks()
reference_masks = reference_masks[:half_y].reshape((half_y*shape[1], shape[2], shape[3]))
masks = generate_masks(
reconstruct_shape=shape[:2],
mask_shape=shape[2:],
dtype=dtype,
lamb=lambda_e,
dpix=dpix,
semiconv=semiconv,
semiconv_pix=semiconv_pix,
cy=cy,
cx=cx,
transformation=identity(),
method="subpix"
).todense()
assert reference_masks.shape == masks.shape
print(reference_masks)
print(masks)
print(reference_masks - masks)
print("maximum difference: ", np.max(np.abs(reference_masks - masks)))
print(np.where(reference_masks != masks))
assert np.any(reference_masks != 0)
assert np.allclose(reference_masks, masks)
def test_ssb_rotate(n_threads):
lt_ctx = lt.Context(InlineJobExecutor(debug=True, inline_threads=n_threads))
dtype = np.float64
scaling = 4
det = 45
shape = (29, 30, det, det)
# ? shape = np.random.uniform(1, 300, (4,1,))
# The acceleration voltage U in keV
U = 300
lamb = wavelength(U)
# STEM pixel size in m, here 50 STEM pixels on 0.5654 nm
dpix = 0.5654/50*1e-9
# STEM semiconvergence angle in radians
semiconv = 25e-3
# Diameter of the primary beam in the diffraction pattern in pixels
semiconv_pix = 78.6649 / scaling
cy = det // 2
cx = det // 2
input_data = (
np.random.uniform(0, 1, np.prod(shape))
* np.linspace(1.0, 1000.0, num=np.prod(shape))
)
input_data = input_data.astype(np.float64).reshape(shape)
data_90deg = np.zeros_like(input_data)
# Rotate 90 degrees clockwise
for y in range(det):
for x in range(det):
data_90deg[:, :, x, det-1-y] = input_data[:, :, y, x]
udf = SSB_UDF(lamb=lamb, dpix=dpix, semiconv=semiconv, semiconv_pix=semiconv_pix,
dtype=dtype, cy=cy, cx=cx, transformation=rotate_deg(-90.))
dataset = MemoryDataSet(
data=data_90deg, tileshape=(20, shape[2], shape[3]), num_partitions=2, sig_dims=2,
)
result = lt_ctx.run_udf(udf=udf, dataset=dataset)
result_f, _ = reference_ssb(input_data, U=U, dpix=dpix, semiconv=semiconv,
semiconv_pix=semiconv_pix, cy=cy, cx=cx)
assert np.allclose(result['fourier'].data, result_f)
def test_wavelength(dtype, atol):
# Handbook of Physics in Medicine and Biology p.40-3
# Table 40.1 'Electron wavelengths at some acceleration voltages used in TEM'
U_voltages = np.array([100, 200, 300, 400, 1000], dtype=dtype)
wavelength_expected = np.array(
[3.70e-12, 2.51e-12, 1.97e-12, 1.64e-12, 0.87e-12], dtype=dtype)
shape = U_voltages.shape[0]
wavelength_result_ssb = np.zeros((shape), dtype=dtype)
for i in range(0, shape):
wavelength_result_ssb[i] = wavelength(U_voltages[i])
assert np.allclose(wavelength_result_ssb, wavelength_expected, atol=atol)
def test_ssb_roi(n_threads):
lt_ctx = lt.Context(InlineJobExecutor(debug=True, inline_threads=n_threads))
dtype = np.float64
scaling = 4
shape = (29, 30, 189 // scaling, 197 // scaling)
# ? shape = np.random.uniform(1, 300, (4,1,))
# The acceleration voltage U in keV
U = 300
lamb = wavelength(U)
# STEM pixel size in m, here 50 STEM pixels on 0.5654 nm
dpix = 0.5654/50*1e-9
# STEM semiconvergence angle in radians
semiconv = 25e-3
# Diameter of the primary beam in the diffraction pattern in pixels
semiconv_pix = 78.6649 / scaling
cy = 93 // scaling
cx = 97 // scaling
input_data = (
np.random.uniform(0, 1, np.prod(shape))
* np.linspace(1.0, 1000.0, num=np.prod(shape))
)
input_data = input_data.astype(np.float64).reshape(shape)
udf = SSB_UDF(lamb=lamb, dpix=dpix, semiconv=semiconv, semiconv_pix=semiconv_pix,
dtype=dtype, cy=cy, cx=cx)
dataset = MemoryDataSet(
data=input_data, tileshape=(20, shape[2], shape[3]), num_partitions=2, sig_dims=2,
)
roi_1 = np.random.choice([True, False], shape[:2])
roi_2 = np.invert(roi_1)
result_1 = lt_ctx.run_udf(udf=udf, dataset=dataset, roi=roi_1)
result_2 = lt_ctx.run_udf(udf=udf, dataset=dataset, roi=roi_2)
result_f, _ = reference_ssb(input_data, U=U, dpix=dpix, semiconv=semiconv,
semiconv_pix=semiconv_pix, cy=cy, cx=cx)
assert np.allclose(result_1['fourier'].data + result_2['fourier'].data, result_f)
def reference_ssb(data, U, dpix, semiconv, semiconv_pix, cy=None, cx=None):
# 'U' - The acceleration voltage U in keV
# 'dpix' - STEM pixel size in m
# 'semiconv' - STEM semiconvergence angle in radians
# 'semiconv_pix' - Diameter of the primary beam in the diffraction pattern in pixels
dpix = np.array(dpix)
reordered = np.moveaxis(data, (0, 1), (2, 3))
ffts = np.fft.fft2(reordered)
rearranged_ffts = np.moveaxis(ffts, (2, 3), (0, 1))
Nblock = np.array(data.shape[0:2])
Nscatter = np.array(data.shape[2:4])
# electron wavelength in m
lamb = wavelength(U)
# spatial freq. step size in scattering space
d_Kf = np.sin(semiconv)/lamb/semiconv_pix
# spatial freq. step size according to probe raster
d_Qp = 1/dpix/Nblock
result_f = np.zeros(data.shape[:2], dtype=rearranged_ffts.dtype)
masks = np.zeros_like(data)
if cx is None:
cx = data.shape[-1] / 2
if cy is None:
cy = data.shape[-2] / 2
y, x = np.ogrid[0:Nscatter[0], 0:Nscatter[1]]
filter_center = circular(
centerX=cx, centerY=cy,
imageSizeX=Nscatter[1], imageSizeY=Nscatter[0],
radius=semiconv_pix,
antialiased=True
).astype(np.float64)
for row in range(Nblock[0]):
for column in range(Nblock[1]):
qp = np.array((row, column))
flip = qp > Nblock / 2
real_qp = qp.copy()
real_qp[flip] = qp[flip] - Nblock[flip]
sy, sx = real_qp * d_Qp / d_Kf
filter_positive = circular(
centerX=cx+sx, centerY=cy+sy,
imageSizeX=Nscatter[1], imageSizeY=Nscatter[0],
radius=semiconv_pix,
antialiased=True
).astype(np.float64)
filter_negative = circular(
centerX=cx-sx, centerY=cy-sy,
imageSizeX=Nscatter[1], imageSizeY=Nscatter[0],
radius=semiconv_pix,
antialiased=True
).astype(np.float64)
mask_positive = filter_center * filter_positive * (filter_negative == 0)
mask_negative = filter_center * filter_negative * (filter_positive == 0)
non_zero_positive = mask_positive.sum()
non_zero_negative = mask_negative.sum()
f = rearranged_ffts[row, column]
if non_zero_positive >= 1 and non_zero_negative >= 1:
tmp = (
(f * mask_positive).sum() / non_zero_positive
- (f * mask_negative).sum() / non_zero_negative
) / 2
result_f[row, column] = tmp
masks[row, column] = ((mask_positive / non_zero_positive) - (
mask_negative / non_zero_negative)) / 2
assert np.allclose(result_f[row, column], (f*masks[row, column]).sum())
else:
assert non_zero_positive < 1
assert non_zero_negative < 1
result_f[0, 0] = (rearranged_ffts[0, 0] * filter_center).sum() / filter_center.sum()
masks[0, 0] = filter_center / filter_center.sum()
return result_f, masks
def test_validate_ssb(real_params, real_intensity_ds, real_plane_wave,
real_reference_ssb, lt_ctx, method, external_container):
'''
The mask generation methods can produce slightly different masks.
Since SSB strongly suppresses noise, including any features
where real space and diffraction space don't properly align,
slight differences in the mask stack can lead to amplifying errors
if the input data contains no actual features and the signal sums up to nearly zero.
For that reason the correctness of mask generation functions shoud be tested on
simulated data that contains a pronounced signal.
Furthermore, this allows to compare the reconstruction with a "ground truth" phase.
'''
dtype = np.float64
shape = real_intensity_ds.shape
# The acceleration voltage U in keV
U = real_params["U"]
lamb = wavelength(U)
# STEM semiconvergence angle in radians
semiconv = real_params["semiconv"]
# Diameter of the primary beam in the diffraction pattern in pixels
semiconv_pix = real_params["semiconv_pix"]
cy = real_params["cy"]
cx = real_params["cx"]
dpix = real_params["dpix"]
transformation = real_params["transformation"]
if external_container:
masks = generate_masks(
reconstruct_shape=shape[:2],
mask_shape=shape[2:],
dtype=dtype,
lamb=lamb,
dpix=dpix,
semiconv=semiconv,
semiconv_pix=semiconv_pix,
cy=cy,
cx=cx,
transformation=transformation,
method=method,
cutoff=1,
)
mask_container = MaskContainer(
mask_factories=lambda: masks, dtype=masks.dtype,
use_sparse='scipy.sparse.csc', count=masks.shape[0],
)
else:
mask_container = None
udf = SSB_UDF(
lamb=lamb, dpix=dpix, semiconv=semiconv, semiconv_pix=semiconv_pix,
dtype=dtype, cy=cy, cx=cx, mask_container=mask_container, method=method,
cutoff=1,
)
ssb_res = lt_ctx.run_udf(udf=udf, dataset=real_intensity_ds)
result_f, reference_masks = real_reference_ssb
# We apply the amplitude scaling to the raw reference SSB result
reference_ssb_raw = np.fft.ifft2(result_f)
reference_ssb_amp = np.abs(reference_ssb_raw)
reference_ssb_phase = np.angle(reference_ssb_raw)
reference_ssb_res = np.sqrt(reference_ssb_amp) * np.exp(1j*reference_ssb_phase)
ssb_phase = ssb_res['phase'].data
ref_phase = np.angle(real_plane_wave)
ssb_amp = ssb_res['amplitude'].data
ref_amp = np.abs(real_plane_wave)
# The phases are usually shifted by a constant offset
# Looking at Std removes the offset
# TODO the current data is at the limit of SSB reconstruction. Better data should be simulated.
# TODO work towards 100 % correspondence with suitable test dataset
assert np.std(ssb_phase - ref_phase) < 0.1 * np.std(ssb_phase)
# Compare reconstructed amplitude
# We can't use std(amp) since the amplitude is nearly constant over the FOV
print("Max ref: ", np.max(np.abs(ssb_amp - ref_amp)), np.max(np.abs(ref_amp)))
assert np.max(np.abs(ssb_amp - ref_amp)) < 0.1 * np.max(np.abs(ref_amp))
# Make sure the methods are at least reasonably comparable
# TODO work towards 100 % correspondence with suitable test dataset
# TODO make the amplitude of the reconstruction match
ssb_res_complex = ssb_res['complex'].data
print(
"Max between: ",
np.max(np.abs(ssb_res_complex - reference_ssb_res)),
np.max(np.abs(ssb_res_complex))
)
print(
"Std between: ",
np.std(ssb_res_complex - reference_ssb_res),
np.std(ssb_res_complex)
)
diffmax = np.max(np.abs(ssb_res_complex - reference_ssb_res))
difflimit = 0.01*np.max(np.abs(ssb_res_complex))
assert diffmax < difflimit
assert np.std(ssb_res_complex - reference_ssb_res) < 0.01*np.std(ssb_res_complex)