diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index e0fe59eee..98e8c95c7 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -975,6 +975,7 @@ def generate_projected_potential( potential_radius_angstroms=3.0, sigma_image_blur_angstroms=0.1, thickness_angstroms=100, + max_num_proj=200, power_scale=1.0, plot_result=False, figsize=(6, 6), @@ -989,9 +990,6 @@ def generate_projected_potential( """ Generate an image of the projected potential of crystal in real space, using cell tiling, and a lookup table of the atomic potentials. - Note that we round atomic positions to the nearest pixel for speed. - - TODO - fix scattering prefactor so that output units are sensible. Parameters ---------- @@ -1006,6 +1004,9 @@ def generate_projected_potential( thickness_angstroms: float Thickness of the sample in Angstroms. Set thickness_thickness_angstroms = 0 to skip thickness projection. + max_num_proj: int + This value prevents this function from projecting a large number of unit + cells along the beam direction, which could be potentially quite slow. power_scale: float Power law scaling of potentials. Set to 2.0 to approximate Z^2 images. plot_result: bool @@ -1054,52 +1055,22 @@ def generate_projected_potential( # Rotate unit cell into projection direction lat_real = self.lat_real.copy() @ orientation_matrix - # Determine unit cell axes to tile over, by selecting 2/3 with largest in-plane component - inds_tile = np.argsort(np.linalg.norm(lat_real[:, 0:2], axis=1))[1:3] + # Determine unit cell axes to tile over, by selecting 2/3 with smallest out-of-plane component + inds_tile = np.argsort(np.abs(lat_real[:, 2]))[0:2] m_tile = lat_real[inds_tile, :] + # Vector projected along optic axis m_proj = np.squeeze(np.delete(lat_real, inds_tile, axis=0)) - # Thickness - if thickness_angstroms > 0: - num_proj = np.round(thickness_angstroms / np.abs(m_proj[2])).astype("int") - if num_proj > 1: - vec_proj = m_proj[:2] / pixel_size_angstroms - shifts = np.arange(num_proj).astype("float") - shifts -= np.mean(shifts) - x_proj = shifts * vec_proj[0] - y_proj = shifts * vec_proj[1] - else: - num_proj = 1 - else: - num_proj = 1 - # Determine tiling range - if thickness_angstroms > 0: - # include the cell height - dz = m_proj[2] * num_proj * 0.5 - p_corners = np.array( - [ - [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, dz], - [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, dz], - [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, dz], - [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, dz], - [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, -dz], - [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, -dz], - [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, -dz], - [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, -dz], - ] - ) - else: - p_corners = np.array( - [ - [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], - [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], - [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], - [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], - ] - ) - + p_corners = np.array( + [ + [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], + [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], + [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], + [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], + ] + ) ab = np.linalg.lstsq(m_tile[:, :2].T, p_corners[:, :2].T, rcond=None)[0] ab = np.floor(ab) a_range = np.array((np.min(ab[0]) - 1, np.max(ab[0]) + 2)) @@ -1115,31 +1086,17 @@ def generate_projected_potential( abc_atoms[:, inds_tile[0]] += a_ind.ravel() abc_atoms[:, inds_tile[1]] += b_ind.ravel() xyz_atoms_ang = abc_atoms @ lat_real - atoms_ID_all_0 = self.numbers[atoms_ind.ravel()] + atoms_ID_all = self.numbers[atoms_ind.ravel()] # Center atoms on image plane - x0 = xyz_atoms_ang[:, 0] / pixel_size_angstroms + im_size[0] / 2.0 - y0 = xyz_atoms_ang[:, 1] / pixel_size_angstroms + im_size[1] / 2.0 - - # if needed, tile atoms in the projection direction - if num_proj > 1: - x = (x0[:, None] + x_proj[None, :]).ravel() - y = (y0[:, None] + y_proj[None, :]).ravel() - atoms_ID_all = np.tile(atoms_ID_all_0, (num_proj, 1)) - else: - x = x0 - y = y0 - atoms_ID_all = atoms_ID_all_0 - # print(x.shape, y.shape) - - # delete atoms outside the field of view - bound = potential_radius_angstroms / pixel_size_angstroms + x = xyz_atoms_ang[:, 0] / pixel_size_angstroms + im_size[0] / 2.0 + y = xyz_atoms_ang[:, 1] / pixel_size_angstroms + im_size[1] / 2.0 atoms_del = np.logical_or.reduce( ( - x <= -bound, - y <= -bound, - x >= im_size[0] + bound, - y >= im_size[1] + bound, + x <= -potential_radius_angstroms / 2, + y <= -potential_radius_angstroms / 2, + x >= im_size[0] + potential_radius_angstroms / 2, + y >= im_size[1] + potential_radius_angstroms / 2, ) ) x = np.delete(x, atoms_del) @@ -1164,16 +1121,18 @@ def generate_projected_potential( for a0 in range(atoms_ID.shape[0]): atom_sf = single_atom_scatter([atoms_ID[a0]]) atoms_lookup[a0, :, :] = atom_sf.projected_potential(atoms_ID[a0], R_2D) - - # if needed, apply gaussian blurring to each atom - if sigma_image_blur_angstroms > 0: - atoms_lookup[a0, :, :] = gaussian_filter( - atoms_lookup[a0, :, :], - sigma_image_blur_angstroms / pixel_size_angstroms, - mode="nearest", - ) atoms_lookup **= power_scale + # Thickness + if thickness_angstroms > 0: + thickness_proj = thickness_angstroms / m_proj[2] + vec_proj = thickness_proj / pixel_size_angstroms * m_proj[:2] + num_proj = (np.ceil(np.linalg.norm(vec_proj)) + 1).astype("int") + num_proj = np.minimum(num_proj, max_num_proj) + + x_proj = np.linspace(-0.5, 0.5, num_proj) * vec_proj[0] + y_proj = np.linspace(-0.5, 0.5, num_proj) * vec_proj[1] + # initialize potential im_potential = np.zeros(im_size) @@ -1181,41 +1140,68 @@ def generate_projected_potential( for a0 in range(atoms_ID_all.shape[0]): ind = np.argmin(np.abs(atoms_ID - atoms_ID_all[a0])) - x_ind = np.round(x[a0]).astype("int") + R_ind - y_ind = np.round(y[a0]).astype("int") + R_ind - x_sub = np.logical_and( - x_ind >= 0, - x_ind < im_size[0], - ) - y_sub = np.logical_and( - y_ind >= 0, - y_ind < im_size[1], - ) - im_potential[x_ind[x_sub][:, None], y_ind[y_sub][None, :]] += atoms_lookup[ - ind - ][x_sub][:, y_sub] + if thickness_angstroms > 0: + for a1 in range(num_proj): + x_ind = np.round(x[a0] + x_proj[a1]).astype("int") + R_ind + y_ind = np.round(y[a0] + y_proj[a1]).astype("int") + R_ind + x_sub = np.logical_and( + x_ind >= 0, + x_ind < im_size[0], + ) + y_sub = np.logical_and( + y_ind >= 0, + y_ind < im_size[1], + ) + + im_potential[ + x_ind[x_sub][:, None], y_ind[y_sub][None, :] + ] += atoms_lookup[ind][x_sub, :][:, y_sub] + + else: + x_ind = np.round(x[a0]).astype("int") + R_ind + y_ind = np.round(y[a0]).astype("int") + R_ind + x_sub = np.logical_and( + x_ind >= 0, + x_ind < im_size[0], + ) + y_sub = np.logical_and( + y_ind >= 0, + y_ind < im_size[1], + ) + + im_potential[ + x_ind[x_sub][:, None], y_ind[y_sub][None, :] + ] += atoms_lookup[ind][x_sub, :][:, y_sub] if thickness_angstroms > 0: im_potential /= num_proj + # if needed, apply gaussian blurring + if sigma_image_blur_angstroms > 0: + sigma_image_blur = sigma_image_blur_angstroms / pixel_size_angstroms + im_potential = gaussian_filter( + im_potential, + sigma_image_blur, + mode="nearest", + ) + if plot_result: # quick plotting of the result int_vals = np.sort(im_potential.ravel()) int_range = np.array( ( int_vals[np.round(0.02 * int_vals.size).astype("int")], - int_vals[np.round(0.999 * int_vals.size).astype("int")], + int_vals[np.round(0.98 * int_vals.size).astype("int")], ) ) fig, ax = plt.subplots(figsize=figsize) ax.imshow( im_potential, - cmap="gray", + cmap="turbo", vmin=int_range[0], vmax=int_range[1], ) - # ax.scatter(y,x,c='r') # for testing ax.set_axis_off() ax.set_aspect("equal") diff --git a/py4DSTEM/process/polar/polar_fits.py b/py4DSTEM/process/polar/polar_fits.py index 18f059b61..7760726a7 100644 --- a/py4DSTEM/process/polar/polar_fits.py +++ b/py4DSTEM/process/polar/polar_fits.py @@ -126,20 +126,20 @@ def fit_amorphous_ring( sigma2 = (radial_range[1] - radial_range[0]) / 4 coefs = (x0, y0, a, b, t, int0, int12, k_bg, sigma0, sigma1, sigma2) - lb = (0, 0, radial_range[0], radial_range[0], -np.inf, 0, 0, 0, 1, 1, 1) - ub = ( - im.shape[0], - im.shape[1], - radial_range[1], - radial_range[1], - np.inf, - np.inf, - np.inf, - np.inf, - np.inf, - np.inf, - np.inf, - ) + lb = (0, 0, radial_range[0], radial_range[0], -np.inf, 0, 0, 0, 1, 1, 1) + ub = ( + im.shape[0], + im.shape[1], + radial_range[1], + radial_range[1], + np.inf, + np.inf, + np.inf, + np.inf, + np.inf, + np.inf, + np.inf, + ) if show_fit_mask: # show image preview of fitting mask diff --git a/py4DSTEM/process/utils/single_atom_scatter.py b/py4DSTEM/process/utils/single_atom_scatter.py index 54443b68f..90397560f 100644 --- a/py4DSTEM/process/utils/single_atom_scatter.py +++ b/py4DSTEM/process/utils/single_atom_scatter.py @@ -57,8 +57,6 @@ def projected_potential(self, Z, R): me = 9.10938356e-31 # Electron charge in Coulomb qe = 1.60217662e-19 - # Electron charge in V-Angstroms - # qe = 14.4 # Permittivity of vacuum eps_0 = 8.85418782e-12 # Bohr's constant @@ -66,21 +64,16 @@ def projected_potential(self, Z, R): fe = np.zeros_like(R) for i in range(5): + # fe += ai[i] * (2 + bi[i] * gsq) / (1 + bi[i] * gsq) ** 2 pre = 2 * np.pi / bi[i] ** 0.5 fe += (ai[i] / bi[i] ** 1.5) * (kn(0, pre * R) + R * kn(1, pre * R)) - # Scale output units - # kappa = (4*np.pi*eps_0) / (2*np.pi*a_0*qe) - # fe *= 2*np.pi**2 / kappa - # # # kappa = (4*np.pi*eps_0) / (2*np.pi*a_0*me) - - # # kappa = (4*np.pi*eps_0) / (2*np.pi*a_0*me) - # # return fe * 2 * np.pi**2 # / kappa - # # if units == "VA": - # return h**2 / (2 * np.pi * me * qe) * 1e18 * fe - # # elif units == "A": - # # return fe * 2 * np.pi**2 / kappa - return fe + # kappa = (4*np.pi*eps_0) / (2*np.pi*a_0*me) + return fe * 2 * np.pi**2 # / kappa + # if units == "VA": + # return h**2 / (2 * np.pi * me * qe) * 1e18 * fe + # elif units == "A": + # return fe * 2 * np.pi**2 / kappa def get_scattering_factor( self, elements=None, composition=None, q_coords=None, units=None