Commit 073aaafd authored by Mukherjee, Debangshu's avatar Mukherjee, Debangshu
Browse files

Filled out STEM simulations, cleaned bugs in custom detector

parent c6b445c0
......@@ -33,7 +33,6 @@ if not os.getenv("READTHEDOCS"):
"h5py >= 2.7.0",
"dask >= 2.0.0",
"numexpr >= 2.6.5",
"nodejs >= 10.0.0",
],
)
else:
......
......@@ -330,69 +330,74 @@ def aperture_image(data4D, center, radius):
return df_image
def custom_detector(data4D, det_inner, det_outer, det_center=(0, 0), mrad_calib=0):
def ROI_from_image(image, med_val, style="over", showfig=True):
if style == "over":
ROI = np.asarray(image > (med_val * np.median(image)), dtype=np.double)
else:
ROI = np.asarray(image < (med_val * np.median(image)), dtype=np.double)
if showfig:
plt.figure(figsize=(15, 15))
plt.imshow(ROI + st.util.image_normalizer(image), cmap="viridis")
plt.title("ROI overlaid")
ROI = ROI.astype(bool)
return ROI
def custom_detector(data4D, inner, outer=0, center=(0, 0), mrad_calib=0):
"""
Generate an image with a custom annular detector
located anywhere in diffraction space
Return STEM image from detector values
Parameters
----------
data4D: ndarray of shape (4,4)
the first two dimensions are Fourier
space, while the next two dimensions
are real space
center: ndarray of shape (1,2)
Center of the circular aperture
radius: float
Radius of the circular aperture
data4D: ndarray
the first two dimensions are Fourier
space, while the next two dimensions
are real space
inner: float
The inner collection angle in Fourier space
in pixels
outer: float, optional
The inner collection angle in Fourier space
in pixels. Default is 0
center: tuple, optional
The center of the 4D-STEM pattern in Fourier
space. Default is (0, 0)
mrad_calib: float, optional
Calibration of the Fourier space. Default
is 0.
Returns
-------
df_image: ndarray of shape (2,2)
Generated virtual dark field image
from the aperture and 4D data
data_det: ndarray
The STEM image from the detector region chosen
Notes
-----
We generate the aperture first, and then make copies
of the aperture to generate a 4D dataset of the same
size as the 4D data. Then we do an element wise
multiplication of this aperture 4D data with the 4D data
and then sum it along the two Fourier directions.
Based on the inner and outer collection angle the a STEM
image is generated from the 4D-STEM dataset. We assume that
the first two dimensions are the Fourier space, while the
next two dimensions are real space scanning dimensions.
"""
if mrad_calib > 0:
det_inner = det_inner * mrad_calib
det_outer = det_outer * mrad_calib
det_center = np.asarray(det_center) * mrad_calib
det_center = np.asarray(det_center)
center = np.asarray(center)
if center[0] == 0:
center[0] = 0.5 * data4D.shape[0]
if center[1] == 0:
center[1] = 0.5 * data4D.shape[1]
yy, xx = np.mgrid[0 : data4D.shape[0], 0 : data4D.shape[1]]
yy -= 0.5 * data4D.shape[0]
xx -= 0.5 * data4D.shape[1]
yy = yy - det_center[1]
xx = xx - det_center[0]
rr = (yy ** 2) + (xx ** 2)
aperture = np.logical_and((rr <= det_outer), (rr >= det_inner))
apt_copy = np.empty(
(data4D.shape[2], data4D.shape[3]) + aperture.shape, dtype=data4D.dtype
)
apt_copy[:] = aperture
apt_copy = np.transpose(apt_copy, (2, 3, 0, 1))
apt_mult = apt_copy * data4D
df_image = np.sum(np.sum(apt_mult, axis=0), axis=0)
return df_image
def ROI_from_image(image, med_val, style="over", showfig=True):
if style == "over":
ROI = np.asarray(image > (med_val * np.median(image)), dtype=np.double)
else:
ROI = np.asarray(image < (med_val * np.median(image)), dtype=np.double)
if showfig:
plt.figure(figsize=(15, 15))
plt.imshow(ROI + st.util.image_normalizer(image), cmap="viridis")
plt.title("ROI overlaid")
ROI = ROI.astype(bool)
return ROI
if mrad_calib > 0:
inner = inner * mrad_calib
outer = outer * mrad_calib
center = center * mrad_calib
yy = yy * mrad_calib
xx = xx * mrad_calib
yy = yy - center[0]
xx = xx - center[1]
rr = ((yy ** 2) + (xx ** 2)) ** 0.5
if outer == 0:
outer = np.amax(rr)
detector = np.logical_and((rr > inner), (rr < outer))
data_det = np.sum(data_4D[detector, :, :], axis=0)
return data_det
@numba.jit
......
......@@ -343,3 +343,121 @@ def slabbing_2D(miller_dir, no_cells, max_hdist):
yy_thirdpass = yy_secondpass[dist_angles2 > (np.pi / 2)]
vals = np.asarray((yy_thirdpass, xx_thirdpass))
return vals.transpose()
def fwxm(probe2D, psize, x=0.5):
p2 = np.abs(probe2D)
pmax = np.amax(p2)
pr = p2 > (x * pmax)
radius = (np.sum(pr) / np.pi) ** 0.5
return radius * psize
def move_probe_mesh(probe, pos, fmeshy, fmeshx):
image_size = np.asarray(probe.shape)
rel_pos = pos - (image_size / 2)
move_mat = np.exp((fmeshx * rel_pos[1]) + (fmeshy * rel_pos[0]))
moved_probe = np.fft.ifft2(move_mat * np.fft.fft2(probe))
return moved_probe
def diff_to_im(cbed_pattern, detector):
val = np.zeros(detector.shape[-1], dtype=np.float32)
for ii in np.arange(detector.shape[-1]):
val[ii] = np.sum(cbed_pattern[detector[:, :, ii]])
return val
def annular_stem(probe, trans_array, prop, pos, coll_angles, pixel_size, voltage_kv):
ypos = pos[0]
xpos = pos[1]
[xpos, ypos] = np.meshgrid(ypos, xpos)
fcal_y = (
np.linspace((-probe.shape[0] / 2), ((probe.shape[0] / 2) - 1), probe.shape[0])
) / probe.shape[0]
fcal_x = (
np.linspace((-probe.shape[1] / 2), ((probe.shape[1] / 2) - 1), probe.shape[1])
) / probe.shape[1]
[fmesh_x, fmesh_y] = np.meshgrid(fcal_x, fcal_y)
mmesh_x = ((-2) * np.pi * 1j * fmesh_x).astype(np.complex64)
mmesh_y = ((-2) * np.pi * 1j * fmesh_y).astype(np.complex64)
fmesh_r = (((fmesh_x ** 2) + (fmesh_y ** 2)) ** 0.5) / pixel_size
radians_r = fmesh_r * st.sim.wavelength_ang(voltage_kv)
detectors = np.zeros(
(radians_r.shape[0], radians_r.shape[1], coll_angles.shape[0]), dtype=bool
)
for cc in np.arange(coll_angles.shape[0]):
inner_angle = coll_angles[cc, 0]
outer_angle = coll_angles[cc, 1]
detectors[:, :, cc] = np.logical_and(
((inner_angle / 1000) < radians_r), (radians_r < (outer_angle / 1000))
)
ypos_da = da.from_array(ypos)
xpos_da = da.from_array(xpos)
probe_sc = client.scatter(probe, broadcast=True)
prop_sc = client.scatter(prop, broadcast=True)
mmesh_y_sc = client.scatter(mmesh_y, broadcast=True)
mmesh_x_sc = client.scatter(mmesh_x, broadcast=True)
trans_array_sc = client.scatter(trans_array, broadcast=True)
detectors_sc = client.scatter(detectors, broadcast=True)
stem_image = []
with tqdm(total=ypos.shape[0]) as pbar:
for ii in np.arange(ypos.shape[0]):
im_vals = []
for jj in np.arange(ypos.shape[1]):
moved_probe = dask.delayed(move_probe_mesh)(
probe_sc, (ypos_da[ii, jj], xpos_da[ii, jj]), mmesh_y_sc, mmesh_x_sc
)
cbed_pattern = dask.delayed(cbed)(moved_probe, trans_array_sc, prop_sc)
aperture_vals = dask.delayed(diff_to_im)(cbed_pattern, detectors_sc)
im_vals.append(aperture_vals)
im_vals = np.asarray(dask.compute(*im_vals))
stem_image.append(im_vals)
pbar.update(1)
stem_image = np.asarray(stem_image)
return stem_image
def pixellated_stem(probe, trans_array, prop, pos):
ypos = pos[0]
xpos = pos[1]
probeft = np.fft.fft2(probe)
probeft /= np.sum(np.abs(probeft))
probe = np.fft.ifft2(probeft)
[xpos, ypos] = np.meshgrid(ypos, xpos)
fcal_y = (
np.linspace((-probe.shape[0] / 2), ((probe.shape[0] / 2) - 1), probe.shape[0])
) / probe.shape[0]
fcal_x = (
np.linspace((-probe.shape[1] / 2), ((probe.shape[1] / 2) - 1), probe.shape[1])
) / probe.shape[1]
[fmesh_x, fmesh_y] = np.meshgrid(fcal_x, fcal_y)
mmesh_x = ((-2) * np.pi * 1j * fmesh_x).astype(np.complex64)
mmesh_y = ((-2) * np.pi * 1j * fmesh_y).astype(np.complex64)
ypos_da = da.from_array(ypos)
xpos_da = da.from_array(xpos)
probe_sc = client.scatter(probe, broadcast=True)
prop_sc = client.scatter(prop, broadcast=True)
mmesh_y_sc = client.scatter(mmesh_y, broadcast=True)
mmesh_x_sc = client.scatter(mmesh_x, broadcast=True)
trans_array_sc = client.scatter(trans_array, broadcast=True)
stem_4D = []
with tqdm(total=ypos.shape[0]) as pbar:
for ii in np.arange(ypos.shape[0]):
im_vals = []
for jj in np.arange(ypos.shape[1]):
moved_probe = dask.delayed(move_probe_mesh)(
probe_sc, (ypos_da[ii, jj], xpos_da[ii, jj]), mmesh_y_sc, mmesh_x_sc
)
cbd = dask.delayed(cbed)(moved_probe, trans_array_sc, prop_sc)
im_vals.append(cbd)
im_vals = np.asarray(dask.compute(*im_vals))
stem_4D.append(im_vals)
pbar.update(1)
stem_4D = np.asarray(stem_4D)
return stem_4D
......@@ -4,12 +4,10 @@ import dask.array as da
import os
import dask
import dask.distributed as dd
import multiprocessing
import struct
import numpy as np
import h5py
import stemtool as st
import numba
import glob
from collections import OrderedDict
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment