def mask_pair_subpix(cy, cx, sy, sx, filter_center, semiconv_pix, cutoff,
mask_shape):
filter_positive = circular(centerX=cx + sx,
centerY=cy + sy,
imageSizeX=mask_shape[1],
imageSizeY=mask_shape[0],
radius=semiconv_pix,
antialiased=True)
filter_negative = circular(centerX=cx - sx,
centerY=cy - sy,
imageSizeX=mask_shape[1],
imageSizeY=mask_shape[0],
radius=semiconv_pix,
antialiased=True)
mask_positive = filter_center * filter_positive * (filter_negative == 0)
mask_negative = filter_center * filter_negative * (filter_positive == 0)
return mask_positive, mask_negative
def get_mask(self, sig_shape):
return masks.circular(
centerY=sig_shape[0] // 2,
centerX=sig_shape[1] // 2,
imageSizeY=sig_shape[0],
imageSizeX=sig_shape[1],
radius=self.radius,
antialiased=True,
)
def mask_pair_subpix(cy, cx, sy, sx, filter_center, semiconv_pix):
'''
Calculate positive and negative trotter mask for circular illumination
using a method with subpixel capability.
Parameters
----------
cy, cx : float
Position of the optical axis on the detector in px, center of illumination
sy, sx : float
Trotter shift value in px
filter_center : numpy.ndarray
Center illumination, i.e. zero order disk. This has to be circular and match the radius
semiconv_pix. It is just passed as an argument for efficientcy to avoid unnecessary
recalculation.
semiconv_pix : float
Semiconvergence angle in measured in detector pixel, i.e. radius of the zero order disk.
Returns
-------
mask_positive, mask_negative : numpy.ndarray
Positive and negative trotter mask
'''
mask_shape = filter_center.shape
filter_positive = circular(centerX=cx + sx,
centerY=cy + sy,
imageSizeX=mask_shape[1],
imageSizeY=mask_shape[0],
radius=semiconv_pix,
antialiased=True)
filter_negative = circular(centerX=cx - sx,
centerY=cy - sy,
imageSizeX=mask_shape[1],
imageSizeY=mask_shape[0],
radius=semiconv_pix,
antialiased=True)
mask_positive = filter_center * filter_positive * (filter_negative == 0)
mask_negative = filter_center * filter_negative * (filter_positive == 0)
return mask_positive, mask_negative
def test_comparison_mask(default_k2is, default_k2is_raw, local_cluster_ctx, lt_ctx):
default_k2is_raw_ds = local_cluster_ctx.load(
"raw", K2IS_TESTDATA_RAW, dtype="u2", nav_shape=(34, 35), sig_shape=(1860, 2048),
)
udf = ApplyMasksUDF(
mask_factories=[lambda: masks.circular(centerX=1024, centerY=930, radius=465,
imageSizeX=2048, imageSizeY=1860)]
)
r1 = local_cluster_ctx.run_udf(udf=udf, dataset=default_k2is)
r2 = local_cluster_ctx.run_udf(udf=udf, dataset=default_k2is_raw_ds)
assert np.allclose(
r1['intensity'],
r2['intensity'],
)
def get_roi(self):
if "roi" not in self.parameters:
return None
params = self.parameters["roi"]
ny, nx = tuple(self.dataset.shape.nav)
if params["shape"] == "disk":
roi = masks.circular(
params["cx"],
params["cy"],
nx,
ny,
params["r"],
)
else:
raise NotImplementedError("unknown shape %s" % params["shape"])
return roi
def cbed_frame(fy=128,
fx=128,
zero=None,
a=None,
b=None,
indices=None,
radius=4,
all_equal=False,
margin=None):
if zero is None:
zero = (fy // 2, fx // 2)
zero = np.array(zero)
if a is None:
a = (fy // 8, 0)
a = np.array(a)
if b is None:
b = make_cartesian(make_polar(a) - (0, np.pi / 2))
b = np.array(b)
if indices is None:
indices = np.mgrid[-10:11, -10:11]
if margin is None:
margin = radius
indices, peaks = frame_peaks(fy=fy,
fx=fx,
zero=zero,
a=a,
b=b,
r=margin,
indices=indices)
data = np.zeros((1, fy, fx), dtype=np.float32)
dists = np.linalg.norm(peaks - zero, axis=-1)
max_val = max(dists.max() + 1, len(peaks) + 1)
for i, p in enumerate(peaks):
data += m.circular(
centerX=p[1],
centerY=p[0],
imageSizeX=fx,
imageSizeY=fy,
radius=radius,
antialiased=True,
) * (1 if all_equal else max(1, max_val - dists[i] + i))
return (data, indices, peaks)
def test_com_comparison_scipy_2_masked(ds_random, lt_ctx):
analysis = lt_ctx.create_com_analysis(
dataset=ds_random, cx=0, cy=0, mask_radius=8
)
results = lt_ctx.run(analysis)
raw_data_by_frame = ds_random.data.reshape((16 * 16, 16, 16))
field_x, field_y = results.field.raw_data
field_x = field_x.reshape((16 * 16))
field_y = field_y.reshape((16 * 16))
disk_mask = masks.circular(
centerX=0, centerY=0,
imageSizeX=16,
imageSizeY=16,
radius=8,
)
for idx in range(16 * 16):
masked_frame = raw_data_by_frame[idx] * disk_mask
scy, scx = measurements.center_of_mass(masked_frame)
assert np.allclose(scx, field_x[idx])
assert np.allclose(scy, field_y[idx])
def test_sum_roi(hdf5_ds_2, tmpdir_factory, lt_ctx, local_cluster_url):
datadir = tmpdir_factory.mktemp('template_tests')
conn = {'connection': {'type': 'tcp', 'address': local_cluster_url}}
path = hdf5_ds_2.path
dataset = _get_hdf5_params(path)
roi_params = {
"shape": "disk",
"cx": 8,
"cy": 8,
"r": 6
}
analysis = [{
"analysisType": "SUM_FRAMES",
"parameters": {
"roi": roi_params
}
}]
notebook = notebook_generator(conn, dataset, analysis, save=True)
notebook = io.StringIO(notebook.getvalue())
nb = nbformat.read(notebook, as_version=4)
ep = ExecutePreprocessor(timeout=600)
ep.preprocess(nb, {"metadata": {"path": datadir}})
data_path = os.path.join(datadir, 'sum_result.npy')
results = np.load(data_path)
nx, ny = hdf5_ds_2.shape.nav
roi = masks.circular(
centerX=roi_params['cx'],
centerY=roi_params['cy'],
imageSizeX=nx,
imageSizeY=ny,
radius=roi_params['r'],
)
udf = SumUDF()
expected = lt_ctx.run_udf(hdf5_ds_2, udf, roi)
assert np.allclose(
results,
expected['intensity'].raw_data,
)
def _make_com(cx, cy, r, shape):
disk_mask = masks.circular(
centerX=cx,
centerY=cy,
imageSizeX=frame_size[1],
imageSizeY=frame_size[0],
radius=r,
)
return [
lambda: disk_mask,
lambda: masks.gradient_x(
imageSizeX=frame_size[1],
imageSizeY=frame_size[0],
dtype=dtype,
) * (np.ones(frame_size) * disk_mask),
lambda: masks.gradient_y(
imageSizeX=frame_size[1],
imageSizeY=frame_size[0],
dtype=dtype,
) * (np.ones(frame_size) * disk_mask),
]
def test_sum_with_roi(lt_ctx):
data = _mk_random(size=(16, 16, 16, 16), dtype='<u2')
dataset = MemoryDataSet(data=data,
tileshape=(2, 16, 16),
num_partitions=32)
roi = {
"shape": "disk",
"cx": 5,
"cy": 6,
"r": 7,
}
analysis = SumAnalysis(dataset=dataset, parameters={
"roi": roi,
})
results = lt_ctx.run(analysis)
mask = masks.circular(roi["cx"], roi["cy"], 16, 16, roi["r"])
assert mask.shape == (16, 16)
assert mask[0, 0] == 0
assert mask[6, 5] == 1
assert mask.dtype == bool
# applying the mask flattens the first two dimensions, so we
# only sum over axis 0 here:
expected = data[mask, ...].sum(axis=(0, ))
assert expected.shape == (16, 16)
assert results['intensity'].raw_data.shape == (16, 16)
# is not equal to results without mask:
assert not np.allclose(results['intensity'].raw_data,
data.sum(axis=(0, 1)))
# ... but rather like `expected`:
assert np.allclose(results['intensity'].raw_data, expected)
assert np.allclose(results['intensity_lin'].raw_data, expected)
def generate_masks(reconstruct_shape,
mask_shape,
dtype,
lamb,
dpix,
semiconv,
semiconv_pix,
transformation=None,
cy=None,
cx=None,
cutoff=1,
cutoff_freq=np.float32('inf'),
method='subpix'):
'''
Generate the trotter mask stack.
The y dimension is trimmed to size(y)//2 + 1 to exploit the inherent
symmetry of the mask stack.
Parameters
----------
reconstruct_shape : tuple(int)
Shape of the reconstructed area
mask_shape : tuple(int)
Shape of the detector
dtype : numpy dtype
dtype to use for the mask stack
lamb : float
Wavelength of the illuminating radiation in m
dpix : float or (float, float)
Scan step in m. Tuple (y, x) in case scan step is different in x and y direction.
semiconv : float
Semiconvergence angle of the illumination in radians
semiconv_pix : float
Semiconvergence angle in measured in detector pixel, i.e. radius of the zero order disk.
transformation : numpy.ndarray, optional
Matrix for affine transformation from the scan coordinate directions
to the detector coordinate directions. This does not include the scale, which is handled by
dpix, lamb, semiconv and semiconv_pix. It should only be used to rotate and flip
the coordinate system as necessary. See also
https://github.com/LiberTEM/LiberTEM/blob/master/src/libertem/corrections/coordinates.py
cy, cx : float, optional
Position of the optical axis on the detector in px, center of illumination.
Default: Center of the detector
cutoff : int, optional
Minimum number of pixels in the positive and negative trotter. This can be used to purge
very small trotters to reduce noise. Default is 1, i.e. no cutoff unless one trotter is
empty.
cutoff_freq: float
Trotters belonging to a spatial frequency higher than this value in reciprocal pixel
coordinates will be cut off.
method : str, optional
Can be :code:`'subpix'`(default) or :code:`'shift'` to switch between
:meth:`mask_pair_subpix` and :meth:`mask_pair_shift` to generate a trotter pair.
Returns
-------
masks : sparse.COO
Masks in sparse.pydata.org COO format. y and x frequency index are FFT shifted, i.e. the
zero frequency is at (0,0) and negative frequencies are in the far quadrant and reversed.
The y frequency index is cut in half with size(y)//2 + 1 to exploit the inherent symmetry
of a real-valued Fourier transform. The y and x index are then flattened to make it
suitable for using it with MaskContainer.
'''
reconstruct_shape = np.array(reconstruct_shape)
dpix = np.array(dpix)
d_Kf = np.sin(semiconv) / lamb / semiconv_pix
d_Qp = 1 / dpix / reconstruct_shape
if cy is None:
cy = mask_shape[0] / 2
if cx is None:
cx = mask_shape[1] / 2
if transformation is None:
transformation = identity()
filter_center = circular(centerX=cx,
centerY=cy,
imageSizeX=mask_shape[1],
imageSizeY=mask_shape[0],
radius=semiconv_pix,
antialiased=True)
half_reconstruct = (reconstruct_shape[0] // 2 + 1, reconstruct_shape[1])
masks = []
for row in range(half_reconstruct[0]):
for column in range(half_reconstruct[1]):
# Do an fftshift of q and p
qp = np.array((row, column))
flip = qp > (reconstruct_shape / 2)
real_qp = qp.copy()
real_qp[flip] = qp[flip] - reconstruct_shape[flip]
if np.sum(real_qp**2) > cutoff_freq**2:
masks.append(empty_mask(mask_shape, dtype=dtype))
continue
# Shift of diffraction order relative to zero order
# without rotation in physical coordinates
real_sy_phys, real_sx_phys = real_qp * d_Qp
# We apply the transformation backwards to go
# from physical orientation to detector orientation,
# while the forward direction in center of mass analysis
# goes from detector coordinates to physical coordinates
# Afterwards, we transform from physical detector coordinates
# to pixel coordinates
sy, sx = ((real_sy_phys, real_sx_phys) @ transformation) / d_Kf
masks.append(
generate_mask(
cy=cy,
cx=cx,
sy=sy,
sx=sx,
filter_center=filter_center,
semiconv_pix=semiconv_pix,
cutoff=cutoff,
dtype=dtype,
method=method,
))
# Since we go through masks in order, this gives a mask stack with
# flattened (q, p) dimension to work with dot product and mask container
masks = sparse.stack(masks)
return masks
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 generate_masks(reconstruct_shape,
mask_shape,
dtype,
lamb,
dpix,
semiconv,
semiconv_pix,
transformation=None,
center=None,
cutoff=1,
cutoff_freq=np.float32('inf'),
method='subpix'):
reconstruct_shape = np.array(reconstruct_shape)
dpix = np.array(dpix)
d_Kf = np.sin(semiconv) / lamb / semiconv_pix
d_Qp = 1 / dpix / reconstruct_shape
if center is None:
center = np.array(mask_shape) / 2
if transformation is None:
transformation = identity()
cy, cx = center
filter_center = circular(centerX=cx,
centerY=cy,
imageSizeX=mask_shape[1],
imageSizeY=mask_shape[0],
radius=semiconv_pix,
antialiased=True)
half_reconstruct = (reconstruct_shape[0] // 2 + 1, reconstruct_shape[1])
masks = []
for row in range(half_reconstruct[0]):
for column in range(half_reconstruct[1]):
# Do an fftshift of q and p
qp = np.array((row, column))
flip = qp > (reconstruct_shape / 2)
real_qp = qp.copy()
real_qp[flip] = qp[flip] - reconstruct_shape[flip]
if np.sum(real_qp**2) > cutoff_freq**2:
masks.append(empty_mask(mask_shape, dtype=dtype))
continue
# Shift of diffraction order relative to zero order
# without rotation in physical coordinates
real_sy_phys, real_sx_phys = real_qp * d_Qp
# We apply the transformation backwards to go
# from physical orientation to detector orientation,
# while the forward direction in center of mass analysis
# goes from detector coordinates to physical coordinates
# Afterwards, we transform from physical detector coordinates
# to pixel coordinates
sy, sx = ((real_sy_phys, real_sx_phys) @ transformation) / d_Kf
masks.append(
generate_mask(
cy=cy,
cx=cx,
sy=sy,
sx=sx,
filter_center=filter_center,
semiconv_pix=semiconv_pix,
cutoff=cutoff,
mask_shape=mask_shape,
dtype=dtype,
method=method,
))
# Since we go through masks in order, this gives a mask stack with
# flattened (q, p) dimension to work with dot product and mask container
masks = sparse.stack(masks)
return masks
def get_task_data(self):
# shorthand, cupy or numpy
xp = self.xp
# Hack to pass a fixed external container
# In particular useful for single-process live processing
# or inline executor
ds_nav = tuple(self.meta.dataset_shape.nav)
y_positions, x_positions = np.mgrid[0:ds_nav[0], 0:ds_nav[1]]
# Precalculated values for Fourier transform
row_steps = -2j*np.pi*np.linspace(0, 1, self.reconstruct_shape[0], endpoint=False)
col_steps = -2j*np.pi*np.linspace(0, 1, self.reconstruct_shape[1], endpoint=False)
if self.meta.roi is None:
y_map = y_positions.flatten()
x_map = x_positions.flatten()
else:
y_map = y_positions[self.meta.roi]
x_map = x_positions[self.meta.roi]
if self.params.filter_center is None:
cy, cx = self.params.center
mask_shape = tuple(self.meta.dataset_shape.sig)
filter_center = circular(
centerX=cx, centerY=cy,
imageSizeX=mask_shape[1], imageSizeY=mask_shape[0],
radius=self.params.semiconv_pix,
antialiased=True
).astype(self.params.dtype)
else:
filter_center = self.params.filter_center.astype(self.params.dtype)
steps_dtype = np.result_type(np.complex64, self.params.dtype)
masks = generate_masks(
reconstruct_shape=self.reconstruct_shape,
mask_shape=tuple(self.meta.dataset_shape.sig),
dtype=self.params.dtype,
wavelength=wavelength(self.params.U),
dpix=self.params.dpix,
semiconv=self.params.semiconv,
semiconv_pix=self.params.semiconv_pix,
center=self.params.center,
transformation=self.params.transformation,
cutoff=self.params.cutoff,
filter_center=filter_center
)
skyline = generate_skyline(
reconstruct_shape=self.reconstruct_shape,
mask_shape=tuple(self.meta.dataset_shape.sig),
dtype=self.params.dtype,
wavelength=wavelength(self.params.U),
dpix=self.params.dpix,
semiconv=self.params.semiconv,
semiconv_pix=self.params.semiconv_pix,
tiling_scheme=self.meta.tiling_scheme,
filter_center=filter_center,
center=self.params.center,
transformation=self.params.transformation,
cutoff=self.params.cutoff,
debug_masks=masks.reshape((
self.reconstruct_shape[0]//2 + 1,
self.reconstruct_shape[1],
*tuple(self.meta.dataset_shape.sig)
)).todense()
)
container = MaskContainer(
mask_factories=lambda: masks, dtype=masks.dtype,
use_sparse='scipy.sparse.csc', count=masks.shape[0], backend=self.meta.backend
)
return {
# Frame positions in the dataset masked by ROI
# to easily access position in dataset when
# processing with ROI applied
"skyline": skyline,
"masks": container,
"filter_center": xp.array(filter_center),
"y_map": xp.array(y_map),
"x_map": xp.array(x_map),
"row_steps": xp.array(row_steps.astype(steps_dtype)),
"col_steps": xp.array(col_steps.astype(steps_dtype)),
}
