def test_dtypes(self, dtype):
arr = np.zeros((3, 3), dtype=dtype)
mask = afr._make_mask_circle_centre(arr, 1)
np.testing.assert_equal(
mask,
np.array([[True, False, True], [False, False, False],
[True, False, True]]))
# Get data
sublattice = am.dummy_data.get_scanning_distortion_sublattice()
# Plotting sublattice
signal = sublattice.signal
signal.plot()
signal._plot.signal_plot.figure.savefig(os.path.join(my_path, 'distortion_signal.png'))
##########################
# Getting cropped atom
atom_i = 200
atom = sublattice.atom_list[atom_i]
atom_x, atom_y = atom.pixel_x, atom.pixel_y
radius = 6
atom_image = atom._get_image_slice_around_atom(sublattice.image, radius)[0]
atom_mask = afr._make_mask_circle_centre(atom_image, radius)
atom_image[atom_mask] = 0
# Getting line center of mass
line_x_com_list = []
for ix in range(2, atom_image.shape[1] - 2):
line_mask_x = atom_mask[:, ix]
com_x_offset = line_mask_x[:round(len(line_mask_x)/2)].sum()
line_x = atom_image[:, ix][np.invert(line_mask_x)]
line_x_com = center_of_mass(line_x)[0] + com_x_offset
line_x_com_list.append(line_x_com)
line_x_com_range = range(len(line_x_com_list))
line_x_com_poly = np.polyfit(line_x_com_range, line_x_com_list, deg=1)
line_x_com_fit = np.poly1d(line_x_com_poly)(line_x_com_range)
line_x_variation = np.array(line_x_com_list) - np.array(line_x_com_fit)
def estimate_local_scanning_distortion(
self, image_data, radius=6, edge_skip=2):
"""Get the amount of local scanning distortion from an atomic column.
This is done by assuming the atomic column has a symmetrical shape,
like Gaussian or Lorentzian. The distortion is calculated by getting
the image_data around the atom, given by the radius parameter.
This gives a square cropped version of the image_data, where the
region outside the radius is masked. For each line in the horizontal
direction, the center of mass is found. This gives a list of
horizontal positions as a function of the vertical lines.
To remove the effects like astigmatism and mistilt a linear fit is
fitted to this list of horizontal positions. This fit is then
subtracted from the horizontal position. The distortion for the
vertical lines is then calculated by getting the standard deviation
of this list of values.
Getting the horizontal distortions is calculated in a similar fashion,
but with a list of vertical positions as a function of the horizontal
lines.
Parameters
----------
image_data : 2D NumPy array
radius : int
Radius of the masked and cropped image. Default 6.
edge_skip : int
When the cropped image is masked with a circle,
the edges will consist of very few unmasked pixels.
The center of mass of these lines will be unreliable, so they're
by default skipped. edge_skip = 2 means the two lines closest to
the edge of the cropped image will be skipped. Default 2.
Returns
-------
scanning_distortion : tuple
Both horizontal and vertical directions. For standard raster scans,
horizontal will be the fast scan direction,
and y the slow scan direction. Thus the returned values will be:
(fast scan, slow scan), where typically the slow scan direction has
the largest amount of distortion.
Examples
--------
>>> sl = am.dummy_data.get_scanning_distortion_sublattice()
>>> atom = sl.atom_list[50]
>>> distortion = atom.estimate_local_scanning_distortion(sl.image)
"""
atom_image = self._get_image_slice_around_atom(image_data, radius)[0]
atom_mask = afr._make_mask_circle_centre(atom_image, radius)
atom_image[atom_mask] = 0
line_x_com_list = []
for ix in range(edge_skip, atom_image.shape[1] - edge_skip):
line_mask_x = atom_mask[:, ix]
com_x_offset = line_mask_x[:round(len(line_mask_x)/2)].sum()
line_x = atom_image[:, ix][np.invert(line_mask_x)]
if np.any(line_x):
line_x_com = center_of_mass(line_x)[0] + com_x_offset
line_x_com_list.append(line_x_com)
line_x_com_range = range(len(line_x_com_list))
line_x_com_poly = np.polyfit(line_x_com_range, line_x_com_list, deg=1)
line_x_com_fit = np.poly1d(line_x_com_poly)(line_x_com_range)
line_x_variation = np.array(line_x_com_list) - np.array(line_x_com_fit)
line_y_com_list = []
for iy in range(edge_skip, atom_image.shape[0] - edge_skip):
line_mask_y = atom_mask[iy]
com_y_offset = line_mask_y[:round(len(line_mask_y)/2)].sum()
line_y = atom_image[iy][np.invert(line_mask_y)]
if np.any(line_y):
line_y_com = center_of_mass(line_y)[0] + com_y_offset
line_y_com_list.append(line_y_com)
line_y_com_range = range(len(line_y_com_list))
line_y_com_poly = np.polyfit(line_y_com_range, line_y_com_list, deg=1)
line_y_com_fit = np.poly1d(line_y_com_poly)(line_y_com_range)
line_y_variation = np.array(line_y_com_list) - np.array(line_y_com_fit)
line_x_std = np.std(line_x_variation)
line_y_std = np.std(line_y_variation)
return line_x_std, line_y_std
def test_radius_2(self):
arr = np.zeros((3, 3))
mask = afr._make_mask_circle_centre(arr, 2)
np.testing.assert_equal(mask == 0, True)
def test_wrong_arr_dimensions(self):
arr = np.zeros((3, 3, 4))
with pytest.raises(ValueError):
afr._make_mask_circle_centre(arr, 2)
def test_radius_2(self):
arr = np.zeros((3, 3))
mask = afr._make_mask_circle_centre(arr, 2)
np.testing.assert_array_equal(mask, np.zeros((3, 3), dtype=np.bool))
def test_different_arr_shapes(self, shape):
arr = np.zeros(shape)
mask = afr._make_mask_circle_centre(arr, 2)
assert arr.shape == mask.shape