def test_ssb_rotate():
ctx = lt.Context(executor=InlineJobExecutor())
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 = 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['pixels'].data, result_f)
def test_ssb():
ctx = lt.Context(executor=InlineJobExecutor())
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
# 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)
LG = np.linspace(1.0,
1000.0,
num=shape[0] * shape[1] * shape[2] * shape[3])
LG = LG.reshape(shape[0], shape[1], shape[2], shape[3])
input_data = input_data * LG
input_data = input_data.astype(np.float64)
udf = SSB_UDF(U=U,
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,
)
result = 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)
# atol = np.max(np.abs(result_f))*0.009
# print(np.max(np.abs(np.abs(result['pixels']) - np.abs(result_f))))
assert np.allclose(np.abs(result['pixels']), np.abs(result_f))
def test_ssb_roi():
ctx = lt.Context(executor=InlineJobExecutor())
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 = ctx.run_udf(udf=udf, dataset=dataset, roi=roi_1)
result_2 = 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['pixels'].data + result_2['pixels'].data,
result_f)
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,
)
result = lt_ctx.run_udf(udf=udf, dataset=real_intensity_ds)
result_f, reference_masks = real_reference_ssb
ssb_res = get_results(result)
# 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 = np.angle(ssb_res)
ref_phase = np.angle(real_plane_wave)
ssb_amp = np.abs(ssb_res)
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
print("Max between: ", np.max(np.abs(ssb_res - reference_ssb_res)),
np.max(np.abs(ssb_res)))
print("Std between: ", np.std(ssb_res - reference_ssb_res),
np.std(ssb_res))
assert np.max(
np.abs(ssb_res - reference_ssb_res)) < 0.01 * np.max(np.abs(ssb_res))
assert np.std(ssb_res - reference_ssb_res) < 0.01 * np.std(ssb_res)
def test_ssb_container(dpix, lt_ctx, backend):
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)
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,
method='subpix')
mask_container = MaskContainer(
mask_factories=lambda: masks,
dtype=masks.dtype,
use_sparse='scipy.sparse.csc',
count=masks.shape[0],
)
udf = SSB_UDF(lamb=lamb,
dpix=dpix,
semiconv=semiconv,
semiconv_pix=semiconv_pix,
dtype=dtype,
cy=cy,
cx=cx,
mask_container=mask_container)
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['pixels'].data - result_f)))
assert np.allclose(result['pixels'].data, result_f)
finally:
if backend == 'cupy':
set_use_cpu(0)