From a15b5d949e46dde97a2ce79e5313db1e8a98284a Mon Sep 17 00:00:00 2001 From: cophus Date: Tue, 14 Nov 2023 10:15:09 -0800 Subject: [PATCH 01/13] initial commit for projected potential --- py4DSTEM/process/diffraction/crystal.py | 115 ++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index fb2911992..2d666c3c3 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -876,6 +876,121 @@ def generate_ring_pattern( if return_calc is True: return radii_unique, intensity_unique + + + + def generate_projected_potential( + self, + im_size = (256,256), + pixel_size_Ang = 0.1, + orientation: Optional[Orientation] = None, + ind_orientation: Optional[int] = 0, + orientation_matrix: Optional[np.ndarray] = None, + zone_axis_lattice: Optional[np.ndarray] = None, + proj_x_lattice: Optional[np.ndarray] = None, + zone_axis_cartesian: Optional[np.ndarray] = None, + proj_x_cartesian: Optional[np.ndarray] = None, + ): + """ + Generate a single diffraction pattern, return all peaks as a pointlist. + + Parameters + ---------- + edge_blend: int, optional + Pixels to blend image at the border + + Returns + -------- + orientation (Orientation): an Orientation class object + + Returns: + im_potential: (np.array) + + """ + + # Determine image size in Angstroms + im_size = np.array(im_size) + im_size_Ang = im_size * pixel_size_Ang + + # Parse orientation inputs + if orientation is not None: + if ind_orientation is None: + orientation_matrix = orientation.matrix[0] + else: + orientation_matrix = orientation.matrix[ind_orientation] + elif orientation_matrix is None: + orientation_matrix = self.parse_orientation( + zone_axis_lattice, proj_x_lattice, zone_axis_cartesian, proj_x_cartesian + ) + # projection directions of potential image + proj_x = orientation_matrix[:,0] \ + / np.linalg.norm(orientation_matrix[:,0]) + proj_y = orientation_matrix[:,1] \ + / np.linalg.norm(orientation_matrix[:,1]) + proj_z = orientation_matrix[:,2] \ + / np.linalg.norm(orientation_matrix[:,2]) + + # Determine unit cell axes to tile over + uvw = self.lat_real / \ + np.linalg.norm(self.lat_real, axis = 1) + test = np.abs(uvw @ proj_z) + inds_tile = np.argsort(test)[:2] + m_tile = self.lat_real[inds_tile,:] + + # Determine tiling range + p = 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.floor(np.linalg.lstsq( + m_tile.T, + p.T, + rcond=None)[0]) + a_range = np.array((np.min(ab[0]),np.max(ab[0]))) + b_range = np.array((np.min(ab[1]),np.max(ab[1]))) + + # Tile unit cell + a_ind, b_ind, atoms_ind = np.meshgrid( + np.arange(a_range[0],a_range[1]), + np.arange(b_range[0],b_range[1]), + np.arange(self.positions.shape[0]), + ) + abc_atoms = self.positions[atoms_ind.ravel(),:] + abc_atoms[:,inds_tile[0]] += a_ind.ravel() + abc_atoms[:,inds_tile[1]] += b_ind.ravel() + xyz_atoms_ang = abc_atoms @ self.lat_real.T + # xyz_atoms_pixels = xyz_atoms_ang / pixel_size_Ang \ + # + im_size/2.0 + + # Project into projected potential image plane + x = proj_x + + # # Lookup table for atomic projected potentials + + # # initialize + im_potential = np.zeros(im_size) + # # im_potential[10:20,10:20] = 1 + + + # Add atoms to potential + + + # test plotting + fig,ax = plt.subplots(figsize = (6,6)) + ax.imshow( + im_potential, + cmap = 'gray', + ) + ax.set_axis_off() + + return im_potential + + + + # Vector conversions and other utilities for Crystal classes def cartesian_to_lattice(self, vec_cartesian): vec_lattice = self.lat_inv @ vec_cartesian From 5b9080d55a11fde3a353979042d48d12c3ff4242 Mon Sep 17 00:00:00 2001 From: cophus Date: Tue, 14 Nov 2023 10:32:33 -0800 Subject: [PATCH 02/13] Atom coordinates mostly working Some tiling issue for image corners --- py4DSTEM/process/diffraction/crystal.py | 35 +++++++++++++++++-------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 2d666c3c3..3a661583e 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -938,19 +938,20 @@ def generate_projected_potential( m_tile = self.lat_real[inds_tile,:] # Determine tiling range - p = np.array([ + 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.floor(np.linalg.lstsq( + p_corners_proj = p_corners @ \ + np.linalg.inv(np.vstack((proj_x, proj_y, proj_z))) + ab = np.round(np.linalg.lstsq( m_tile.T, - p.T, + p_corners_proj.T, rcond=None)[0]) - a_range = np.array((np.min(ab[0]),np.max(ab[0]))) - b_range = np.array((np.min(ab[1]),np.max(ab[1]))) + a_range = np.array((np.min(ab[0])-1,np.max(ab[0])+1)) + b_range = np.array((np.min(ab[1])-1,np.max(ab[1])+1)) # Tile unit cell a_ind, b_ind, atoms_ind = np.meshgrid( @@ -962,15 +963,25 @@ 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 @ self.lat_real.T - # xyz_atoms_pixels = xyz_atoms_ang / pixel_size_Ang \ - # + im_size/2.0 # Project into projected potential image plane - x = proj_x + x = (xyz_atoms_ang @ proj_x) / pixel_size_Ang + im_size[0]/2.0 + y = (xyz_atoms_ang @ proj_y) / pixel_size_Ang + im_size[1]/2.0 + atoms_del = np.logical_or.reduce(( + x < 0, + y < 0, + x > im_size[0], + y > im_size[1], + )) + x = np.delete(x, atoms_del) + y = np.delete(y, atoms_del) + + # Lookup table for atomic projected potentials + + - # # Lookup table for atomic projected potentials - # # initialize + # initialize potential im_potential = np.zeros(im_size) # # im_potential[10:20,10:20] = 1 @@ -984,7 +995,9 @@ def generate_projected_potential( im_potential, cmap = 'gray', ) + ax.scatter(y,x) ax.set_axis_off() + ax.set_aspect('equal') return im_potential From 36ddbcc656c9d1fc006bcc1759ec7290d2c27389 Mon Sep 17 00:00:00 2001 From: cophus Date: Tue, 14 Nov 2023 11:12:11 -0800 Subject: [PATCH 03/13] Projected potential now working with bugs projection algebra definitely has a bug --- py4DSTEM/process/diffraction/crystal.py | 73 ++++++++++++++----- py4DSTEM/process/utils/single_atom_scatter.py | 31 +++++++- 2 files changed, 84 insertions(+), 20 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 3a661583e..eb8642b25 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -1,6 +1,7 @@ # Functions for calculating diffraction patterns, matching them to experiments, and creating orientation and phase maps. import numpy as np +from scipy.ndimage import gaussian_filter import matplotlib.pyplot as plt from matplotlib.patches import Circle from fractions import Fraction @@ -883,6 +884,9 @@ def generate_projected_potential( self, im_size = (256,256), pixel_size_Ang = 0.1, + potential_radius_Ang = 3.0, + sigma_image_blur_Ang = 0.1, + plot_result = False, orientation: Optional[Orientation] = None, ind_orientation: Optional[int] = 0, orientation_matrix: Optional[np.ndarray] = None, @@ -963,6 +967,7 @@ 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 @ self.lat_real.T + atoms_ID_all = self.numbers[atoms_ind.ravel()] # Project into projected potential image plane x = (xyz_atoms_ang @ proj_x) / pixel_size_Ang + im_size[0]/2.0 @@ -975,35 +980,65 @@ def generate_projected_potential( )) x = np.delete(x, atoms_del) y = np.delete(y, atoms_del) + atoms_ID_all = np.delete(atoms_ID_all, atoms_del) - # Lookup table for atomic projected potentials - - + # Coordinate system for atomic projected potentials + potential_radius = np.ceil(potential_radius_Ang / pixel_size_Ang) + R = np.arange(0.5-potential_radius,potential_radius+0.5) + R_ind = R.astype('int') + R_2D = np.sqrt(R[:,None]**2 + R[None,:]**2) + # Lookup table for atomic projected potentials + atoms_ID = np.unique(self.numbers) + atoms_lookup = np.zeros(( + atoms_ID.shape[0], + R_2D.shape[0], + R_2D.shape[1], + )) + 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) # initialize potential im_potential = np.zeros(im_size) - # # im_potential[10:20,10:20] = 1 - - - # Add atoms to potential + # Add atoms to potential image + for a0 in range(atoms_ID_all.shape[0]): + ind = np.argmin(np.abs(atoms_ID - atoms_ID_all[a0])) + + x_ind = np.clip( + np.round(x[a0]).astype('int') + R_ind, + 0, + im_size[0]-1) + y_ind = np.clip( + np.round(y[a0]).astype('int') + R_ind, + 0, + im_size[1]-1) + + im_potential[x_ind[None,:],y_ind[:,None]] += atoms_lookup[ind] + + # if needed, apply gaussian blurring + if sigma_image_blur_Ang > 0: + sigma_image_blur = sigma_image_blur_Ang / pixel_size_Ang + im_potential = gaussian_filter( + im_potential, + sigma_image_blur, + mode = 'nearest', + ) - # test plotting - fig,ax = plt.subplots(figsize = (6,6)) - ax.imshow( - im_potential, - cmap = 'gray', - ) - ax.scatter(y,x) - ax.set_axis_off() - ax.set_aspect('equal') + if plot_result: + # test plotting + fig,ax = plt.subplots(figsize = (6,6)) + ax.imshow( + im_potential, + cmap = 'gray', + ) + # ax.scatter(y,x) + ax.set_axis_off() + ax.set_aspect('equal') return im_potential - - - # Vector conversions and other utilities for Crystal classes def cartesian_to_lattice(self, vec_cartesian): vec_lattice = self.lat_inv @ vec_cartesian diff --git a/py4DSTEM/process/utils/single_atom_scatter.py b/py4DSTEM/process/utils/single_atom_scatter.py index 8d6e2a891..9085afd93 100644 --- a/py4DSTEM/process/utils/single_atom_scatter.py +++ b/py4DSTEM/process/utils/single_atom_scatter.py @@ -1,6 +1,6 @@ import numpy as np import os - +from scipy.special import kn class single_atom_scatter(object): """ @@ -46,6 +46,35 @@ def electron_scattering_factor(self, Z, gsq, units="A"): elif units == "A": return fe + def projected_potential(self, Z, R): + ai = self.e_scattering_factors[Z - 1, 0:10:2] + bi = self.e_scattering_factors[Z - 1, 1:10:2] + + # Planck's constant in Js + h = 6.62607004e-34 + # Electron rest mass in kg + me = 9.10938356e-31 + # Electron charge in Coulomb + qe = 1.60217662e-19 + # Permittivity of vacuum + eps_0 = 8.85418782e-12 + # Bohr's constant + a_0 = 5.29177210903e-11 + + 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)) + + # 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 ): From a0d0f15e60618c962a0667801cdf1ccdc20312d4 Mon Sep 17 00:00:00 2001 From: cophus Date: Tue, 14 Nov 2023 11:14:18 -0800 Subject: [PATCH 04/13] minor tweak --- py4DSTEM/process/diffraction/crystal.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index eb8642b25..603e1a6cb 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -966,7 +966,8 @@ def generate_projected_potential( abc_atoms = self.positions[atoms_ind.ravel(),:] abc_atoms[:,inds_tile[0]] += a_ind.ravel() abc_atoms[:,inds_tile[1]] += b_ind.ravel() - xyz_atoms_ang = abc_atoms @ self.lat_real.T + # NOTE - should this be self.lat_real.T? + xyz_atoms_ang = abc_atoms @ self.lat_real atoms_ID_all = self.numbers[atoms_ind.ravel()] # Project into projected potential image plane From 9b726e55bf8d0581a2116cf083997a3fe4ad0282 Mon Sep 17 00:00:00 2001 From: Colin Date: Fri, 8 Dec 2023 13:49:03 -0800 Subject: [PATCH 05/13] Projected potentials fixed? --- py4DSTEM/process/diffraction/crystal.py | 52 +++++++++++-------------- 1 file changed, 23 insertions(+), 29 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 603e1a6cb..d1cb04401 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -926,36 +926,31 @@ def generate_projected_potential( orientation_matrix = self.parse_orientation( zone_axis_lattice, proj_x_lattice, zone_axis_cartesian, proj_x_cartesian ) - # projection directions of potential image - proj_x = orientation_matrix[:,0] \ - / np.linalg.norm(orientation_matrix[:,0]) - proj_y = orientation_matrix[:,1] \ - / np.linalg.norm(orientation_matrix[:,1]) - proj_z = orientation_matrix[:,2] \ - / np.linalg.norm(orientation_matrix[:,2]) - - # Determine unit cell axes to tile over - uvw = self.lat_real / \ - np.linalg.norm(self.lat_real, axis = 1) - test = np.abs(uvw @ proj_z) - inds_tile = np.argsort(test)[:2] - m_tile = self.lat_real[inds_tile,:] + + # 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] + m_tile = lat_real[inds_tile,:] # Determine tiling range 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], + [ im_size_Ang[0]*0.5, im_size_Ang[1]*0.5, 0.0], ]) - p_corners_proj = p_corners @ \ - np.linalg.inv(np.vstack((proj_x, proj_y, proj_z))) - ab = np.round(np.linalg.lstsq( - m_tile.T, - p_corners_proj.T, - rcond=None)[0]) - a_range = np.array((np.min(ab[0])-1,np.max(ab[0])+1)) - b_range = np.array((np.min(ab[1])-1,np.max(ab[1])+1)) + 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)) + b_range = np.array((np.min(ab[1])-1,np.max(ab[1])+2)) # Tile unit cell a_ind, b_ind, atoms_ind = np.meshgrid( @@ -966,13 +961,12 @@ def generate_projected_potential( abc_atoms = self.positions[atoms_ind.ravel(),:] abc_atoms[:,inds_tile[0]] += a_ind.ravel() abc_atoms[:,inds_tile[1]] += b_ind.ravel() - # NOTE - should this be self.lat_real.T? - xyz_atoms_ang = abc_atoms @ self.lat_real + xyz_atoms_ang = abc_atoms @ lat_real atoms_ID_all = self.numbers[atoms_ind.ravel()] - # Project into projected potential image plane - x = (xyz_atoms_ang @ proj_x) / pixel_size_Ang + im_size[0]/2.0 - y = (xyz_atoms_ang @ proj_y) / pixel_size_Ang + im_size[1]/2.0 + # Center atoms on image plane + x = xyz_atoms_ang[:,0] / pixel_size_Ang + im_size[0]/2.0 + y = xyz_atoms_ang[:,1] / pixel_size_Ang + im_size[1]/2.0 atoms_del = np.logical_or.reduce(( x < 0, y < 0, @@ -1032,7 +1026,7 @@ def generate_projected_potential( fig,ax = plt.subplots(figsize = (6,6)) ax.imshow( im_potential, - cmap = 'gray', + cmap = 'turbo', ) # ax.scatter(y,x) ax.set_axis_off() From a9b05edc7bd344757479ac662db8cd6866d2638f Mon Sep 17 00:00:00 2001 From: Colin Date: Fri, 2 Feb 2024 10:43:53 -0800 Subject: [PATCH 06/13] adding figsize --- py4DSTEM/process/diffraction/crystal.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index d1cb04401..7d2606576 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -887,6 +887,7 @@ def generate_projected_potential( potential_radius_Ang = 3.0, sigma_image_blur_Ang = 0.1, plot_result = False, + figsize = (6,6), orientation: Optional[Orientation] = None, ind_orientation: Optional[int] = 0, orientation_matrix: Optional[np.ndarray] = None, @@ -1023,7 +1024,7 @@ def generate_projected_potential( if plot_result: # test plotting - fig,ax = plt.subplots(figsize = (6,6)) + fig,ax = plt.subplots(figsize = figsize) ax.imshow( im_potential, cmap = 'turbo', From 43a396cf7fd5ef7211e3c9431259b77ebc42c429 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 4 Mar 2024 11:33:34 -0800 Subject: [PATCH 07/13] update docstring --- py4DSTEM/process/diffraction/crystal.py | 39 ++++++++++++++++++++----- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 50ddc3087..6f10d8ab1 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -988,19 +988,44 @@ def generate_projected_potential( proj_x_cartesian: Optional[np.ndarray] = None, ): """ - Generate a single diffraction pattern, return all peaks as a pointlist. + Generate an image of the projected potential of crystal in real space, + using cell tiling, and a lookup table of the atomic potentials. Parameters ---------- - edge_blend: int, optional - Pixels to blend image at the border + im_size: tuple, list, np.array + (2,) vector specifying the output size in pixels. + pixel_size_Ang: float + Pixel size in Angstroms. + potential_radius_Ang: float + Radius in Angstroms for how far to integrate the atomic potentials + sigma_image_blur_Ang: float + Image blurring in Angstroms. + plot_result: bool + Plot the projected potential image. + figsize: + (2,) vector giving the size of the output. + + orientation: Orientation + An Orientation class object + ind_orientation: int + If input is an Orientation class object with multiple orientations, + this input can be used to select a specific orientation. + orientation_matrix: array + (3,3) orientation matrix, where columns represent projection directions. + zone_axis_lattice: array + (3,) projection direction in lattice indices + proj_x_lattice: array) + (3,) x-axis direction in lattice indices + zone_axis_cartesian: array + (3,) cartesian projection direction + proj_x_cartesian: array + (3,) cartesian projection direction Returns -------- - orientation (Orientation): an Orientation class object - - Returns: - im_potential: (np.array) + im_potential: (np.array) + Output image of the projected potential. """ From 16c727f7a84e74bada2c9fc118d4de49aea7c6a0 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 4 Mar 2024 13:16:53 -0800 Subject: [PATCH 08/13] Adding thickness projection --- py4DSTEM/process/diffraction/crystal.py | 56 +++++++++++++++---------- 1 file changed, 35 insertions(+), 21 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 6f10d8ab1..de82f443d 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -974,9 +974,10 @@ def generate_ring_pattern( def generate_projected_potential( self, im_size = (256,256), - pixel_size_Ang = 0.1, - potential_radius_Ang = 3.0, - sigma_image_blur_Ang = 0.1, + pixel_size_angstroms = 0.1, + potential_radius_angstroms = 3.0, + sigma_image_blur_angstroms = 0.1, + thickness_angstroms = 100, plot_result = False, figsize = (6,6), orientation: Optional[Orientation] = None, @@ -995,12 +996,14 @@ def generate_projected_potential( ---------- im_size: tuple, list, np.array (2,) vector specifying the output size in pixels. - pixel_size_Ang: float + pixel_size_angstroms: float Pixel size in Angstroms. - potential_radius_Ang: float + potential_radius_angstroms: float Radius in Angstroms for how far to integrate the atomic potentials - sigma_image_blur_Ang: float + sigma_image_blur_angstroms: float Image blurring in Angstroms. + thickness_angstroms: float + Thickness of the sample in Angstroms. plot_result: bool Plot the projected potential image. figsize: @@ -1031,7 +1034,7 @@ def generate_projected_potential( # Determine image size in Angstroms im_size = np.array(im_size) - im_size_Ang = im_size * pixel_size_Ang + im_size_Ang = im_size * pixel_size_angstroms # Parse orientation inputs if orientation is not None: @@ -1052,6 +1055,8 @@ def generate_projected_potential( np.linalg.norm(lat_real[:,0:2],axis=1) )[1:3] m_tile = lat_real[inds_tile,:] + # Vector projected along zone axis + m_proj = np.squeeze(np.delete(lat_real,inds_tile,axis=0)) # Determine tiling range p_corners = np.array([ @@ -1082,8 +1087,8 @@ def generate_projected_potential( atoms_ID_all = self.numbers[atoms_ind.ravel()] # Center atoms on image plane - x = xyz_atoms_ang[:,0] / pixel_size_Ang + im_size[0]/2.0 - y = xyz_atoms_ang[:,1] / pixel_size_Ang + im_size[1]/2.0 + 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 < 0, y < 0, @@ -1095,7 +1100,7 @@ def generate_projected_potential( atoms_ID_all = np.delete(atoms_ID_all, atoms_del) # Coordinate system for atomic projected potentials - potential_radius = np.ceil(potential_radius_Ang / pixel_size_Ang) + potential_radius = np.ceil(potential_radius_angstroms / pixel_size_angstroms) R = np.arange(0.5-potential_radius,potential_radius+0.5) R_ind = R.astype('int') R_2D = np.sqrt(R[:,None]**2 + R[None,:]**2) @@ -1111,6 +1116,13 @@ def generate_projected_potential( atom_sf = single_atom_scatter([atoms_ID[a0]]) atoms_lookup[a0,:,:] = atom_sf.projected_potential(atoms_ID[a0], R_2D) + # Projected thickness + num_proj = thickness_angstroms / m_proj[2] + vec_proj = num_proj / pixel_size_angstroms * m_proj[:2] + print(m_proj) + print(vec_proj) + + # initialize potential im_potential = np.zeros(im_size) @@ -1118,20 +1130,22 @@ 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.clip( - np.round(x[a0]).astype('int') + R_ind, - 0, - im_size[0]-1) - y_ind = np.clip( - np.round(y[a0]).astype('int') + R_ind, - 0, - im_size[1]-1) + 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[None,:],y_ind[:,None]] += atoms_lookup[ind] + im_potential[x_ind[x_sub][:,None],y_ind[y_sub][None,:]] += atoms_lookup[ind][x_sub,:][:,y_sub] # if needed, apply gaussian blurring - if sigma_image_blur_Ang > 0: - sigma_image_blur = sigma_image_blur_Ang / pixel_size_Ang + 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, From a1c83c35044bc55a86f8746594f5b1226ec4e9d6 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 4 Mar 2024 14:24:32 -0800 Subject: [PATCH 09/13] Fourier method makes boundary conditions difficult --- py4DSTEM/process/diffraction/crystal.py | 63 +++++++++++++++++++++---- 1 file changed, 55 insertions(+), 8 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index de82f443d..ed09203ab 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -1055,7 +1055,7 @@ def generate_projected_potential( np.linalg.norm(lat_real[:,0:2],axis=1) )[1:3] m_tile = lat_real[inds_tile,:] - # Vector projected along zone axis + # Vector projected along optic axis m_proj = np.squeeze(np.delete(lat_real,inds_tile,axis=0)) # Determine tiling range @@ -1116,13 +1116,6 @@ def generate_projected_potential( atom_sf = single_atom_scatter([atoms_ID[a0]]) atoms_lookup[a0,:,:] = atom_sf.projected_potential(atoms_ID[a0], R_2D) - # Projected thickness - num_proj = thickness_angstroms / m_proj[2] - vec_proj = num_proj / pixel_size_angstroms * m_proj[:2] - print(m_proj) - print(vec_proj) - - # initialize potential im_potential = np.zeros(im_size) @@ -1142,6 +1135,59 @@ def generate_projected_potential( ) im_potential[x_ind[x_sub][:,None],y_ind[y_sub][None,:]] += atoms_lookup[ind][x_sub,:][:,y_sub] + + # Projected thickness + num_proj = thickness_angstroms / m_proj[2] + vec_proj = num_proj / pixel_size_angstroms * m_proj[:2] + num_points = (np.ceil(np.linalg.norm(vec_proj))*2+1).astype('int') + x = np.linspace(-0.5,0.5,num_points)*vec_proj[0] + im_size[0]/2 + y = np.linspace(-0.5,0.5,num_points)*vec_proj[1] + im_size[1]/2 + xF = np.floor(x).astype('int') + yF = np.floor(y).astype('int') + dx = x - xF + dy = y - yF + im_thickness = np.reshape( + np.bincount( + np.ravel_multi_index((xF ,yF ),im_size,mode='wrap'), + (1-dx)*(1-dy) / num_points, + np.prod(im_size), + ) + \ + np.bincount( + np.ravel_multi_index((xF+1,yF ),im_size,mode='wrap'), + ( dx)*(1-dy) / num_points, + np.prod(im_size), + ) + \ + np.bincount( + np.ravel_multi_index((xF ,yF+1),im_size,mode='wrap'), + (1-dx)*( dy) / num_points, + np.prod(im_size), + ) + \ + np.bincount( + np.ravel_multi_index((xF+1,yF+1),im_size,mode='wrap'), + ( dx)*( dy) / num_points, + np.prod(im_size), + ), + im_size, + ) + # pad images and perform correlation in Fourier space + im_potential = np.real( + np.fft.ifft2( + np.fft.fft2(np.pad(im_potential,((0,im_size[0]),(0,im_size[1])))) \ + * np.fft.fft2(np.pad(im_thickness,((0,im_size[0]),(0,im_size[1])))), + ), + )[im_size[0]//2:im_size[0]+im_size[0]//2,im_size[1]//2:im_size[1]+im_size[1]//2] + + # fig,ax = plt.subplots(figsize=(12,12)) + # ax.imshow( + # np.pad(im_thickness,((0,im_size[0]),(0,im_size[1]))), + # ) + # im_thickness = np.zeros(im_size) + # for a0 in range(num_points): + # x = + + + + # if needed, apply gaussian blurring if sigma_image_blur_angstroms > 0: @@ -1152,6 +1198,7 @@ def generate_projected_potential( mode = 'nearest', ) + if plot_result: # test plotting fig,ax = plt.subplots(figsize = figsize) From c8d661ac49e8c3ba5bfc36a681fe4e3aa34ee5d8 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 4 Mar 2024 15:19:37 -0800 Subject: [PATCH 10/13] Updated plotting --- py4DSTEM/process/diffraction/crystal.py | 123 +++++++++++------------- 1 file changed, 54 insertions(+), 69 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index ed09203ab..8b59c3c92 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -978,6 +978,7 @@ def generate_projected_potential( potential_radius_angstroms = 3.0, sigma_image_blur_angstroms = 0.1, thickness_angstroms = 100, + power_scale = 1.0, plot_result = False, figsize = (6,6), orientation: Optional[Orientation] = None, @@ -1003,7 +1004,10 @@ def generate_projected_potential( sigma_image_blur_angstroms: float Image blurring in Angstroms. thickness_angstroms: float - Thickness of the sample in Angstroms. + Thickness of the sample in Angstroms. + Set thickness_thickness_angstroms = 0 to skip thickness projection. + power_scale: float + Power law scaling of potentials. Set to 2.0 to approximate Z^2 images. plot_result: bool Plot the projected potential image. figsize: @@ -1090,10 +1094,10 @@ def generate_projected_potential( 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 < 0, - y < 0, - x > im_size[0], - y > im_size[1], + 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) y = np.delete(y, atoms_del) @@ -1115,6 +1119,15 @@ 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) + 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') + 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) @@ -1123,71 +1136,37 @@ 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], - ) + 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] + im_potential[x_ind[x_sub][:,None],y_ind[y_sub][None,:]] += atoms_lookup[ind][x_sub,:][:,y_sub] - # Projected thickness - num_proj = thickness_angstroms / m_proj[2] - vec_proj = num_proj / pixel_size_angstroms * m_proj[:2] - num_points = (np.ceil(np.linalg.norm(vec_proj))*2+1).astype('int') - x = np.linspace(-0.5,0.5,num_points)*vec_proj[0] + im_size[0]/2 - y = np.linspace(-0.5,0.5,num_points)*vec_proj[1] + im_size[1]/2 - xF = np.floor(x).astype('int') - yF = np.floor(y).astype('int') - dx = x - xF - dy = y - yF - im_thickness = np.reshape( - np.bincount( - np.ravel_multi_index((xF ,yF ),im_size,mode='wrap'), - (1-dx)*(1-dy) / num_points, - np.prod(im_size), - ) + \ - np.bincount( - np.ravel_multi_index((xF+1,yF ),im_size,mode='wrap'), - ( dx)*(1-dy) / num_points, - np.prod(im_size), - ) + \ - np.bincount( - np.ravel_multi_index((xF ,yF+1),im_size,mode='wrap'), - (1-dx)*( dy) / num_points, - np.prod(im_size), - ) + \ - np.bincount( - np.ravel_multi_index((xF+1,yF+1),im_size,mode='wrap'), - ( dx)*( dy) / num_points, - np.prod(im_size), - ), - im_size, - ) - # pad images and perform correlation in Fourier space - im_potential = np.real( - np.fft.ifft2( - np.fft.fft2(np.pad(im_potential,((0,im_size[0]),(0,im_size[1])))) \ - * np.fft.fft2(np.pad(im_thickness,((0,im_size[0]),(0,im_size[1])))), - ), - )[im_size[0]//2:im_size[0]+im_size[0]//2,im_size[1]//2:im_size[1]+im_size[1]//2] - - # fig,ax = plt.subplots(figsize=(12,12)) - # ax.imshow( - # np.pad(im_thickness,((0,im_size[0]),(0,im_size[1]))), - # ) - # im_thickness = np.zeros(im_size) - # for a0 in range(num_points): - # x = - - + 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: @@ -1198,15 +1177,21 @@ def generate_projected_potential( mode = 'nearest', ) - if plot_result: - # test plotting + # 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.98*int_vals.size).astype('int')], + )) + fig,ax = plt.subplots(figsize = figsize) ax.imshow( im_potential, cmap = 'turbo', + vmin = int_range[0], + vmax = int_range[1], ) - # ax.scatter(y,x) ax.set_axis_off() ax.set_aspect('equal') From d52eaed4be9bac79d3036d3c9f685bd29cb0e9ce Mon Sep 17 00:00:00 2001 From: Colin Date: Wed, 26 Jun 2024 21:09:47 -0700 Subject: [PATCH 11/13] Fix for the infinite projection error --- py4DSTEM/process/diffraction/crystal.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 8b59c3c92..9765d9d22 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -978,6 +978,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), @@ -1006,6 +1007,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 @@ -1055,10 +1059,17 @@ def generate_projected_potential( 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] + # m_tile = lat_real[inds_tile,:] + + # Determine unit cell axes to tile over, by selecting 2/3 with smallest out-of-plane component inds_tile = np.argsort( - np.linalg.norm(lat_real[:,0:2],axis=1) - )[1:3] + 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)) @@ -1126,6 +1137,9 @@ def generate_projected_potential( 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] From 0763f9836ff3a02e50dbbda73b74498e931fb391 Mon Sep 17 00:00:00 2001 From: Colin Date: Wed, 26 Jun 2024 21:10:40 -0700 Subject: [PATCH 12/13] cleaning up --- py4DSTEM/process/diffraction/crystal.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 9765d9d22..81962d010 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -1058,12 +1058,6 @@ 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] - # m_tile = lat_real[inds_tile,:] - # 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]) @@ -1137,7 +1131,6 @@ def generate_projected_potential( 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] From be4d9194ce1ea8987650d3aa62e231b046a37b08 Mon Sep 17 00:00:00 2001 From: Colin Date: Wed, 26 Jun 2024 21:11:53 -0700 Subject: [PATCH 13/13] Black formatting --- py4DSTEM/process/diffraction/crystal.py | 171 +++++++++--------- py4DSTEM/process/utils/single_atom_scatter.py | 10 +- 2 files changed, 92 insertions(+), 89 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 81962d010..98e8c95c7 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -968,20 +968,17 @@ def generate_ring_pattern( if return_calc is True: return radii_unique, intensity_unique - - - def generate_projected_potential( self, - im_size = (256,256), - pixel_size_angstroms = 0.1, - 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), + im_size=(256, 256), + pixel_size_angstroms=0.1, + 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), orientation: Optional[Orientation] = None, ind_orientation: Optional[int] = 0, orientation_matrix: Optional[np.ndarray] = None, @@ -1005,7 +1002,7 @@ def generate_projected_potential( sigma_image_blur_angstroms: float Image blurring in Angstroms. thickness_angstroms: float - Thickness of the sample in Angstroms. + 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 @@ -1016,26 +1013,26 @@ def generate_projected_potential( Plot the projected potential image. figsize: (2,) vector giving the size of the output. - - orientation: Orientation + + orientation: Orientation An Orientation class object ind_orientation: int If input is an Orientation class object with multiple orientations, this input can be used to select a specific orientation. - orientation_matrix: array + orientation_matrix: array (3,3) orientation matrix, where columns represent projection directions. - zone_axis_lattice: array + zone_axis_lattice: array (3,) projection direction in lattice indices - proj_x_lattice: array) + proj_x_lattice: array) (3,) x-axis direction in lattice indices - zone_axis_cartesian: array + zone_axis_cartesian: array (3,) cartesian projection direction - proj_x_cartesian: array + proj_x_cartesian: array (3,) cartesian projection direction Returns -------- - im_potential: (np.array) + im_potential: (np.array) Output image of the projected potential. """ @@ -1059,82 +1056,82 @@ def generate_projected_potential( lat_real = self.lat_real.copy() @ orientation_matrix # 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,:] + 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)) + m_proj = np.squeeze(np.delete(lat_real, inds_tile, axis=0)) # Determine tiling range - 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] + 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)) - b_range = np.array((np.min(ab[1])-1,np.max(ab[1])+2)) + a_range = np.array((np.min(ab[0]) - 1, np.max(ab[0]) + 2)) + b_range = np.array((np.min(ab[1]) - 1, np.max(ab[1]) + 2)) # Tile unit cell a_ind, b_ind, atoms_ind = np.meshgrid( - np.arange(a_range[0],a_range[1]), - np.arange(b_range[0],b_range[1]), + np.arange(a_range[0], a_range[1]), + np.arange(b_range[0], b_range[1]), np.arange(self.positions.shape[0]), ) - abc_atoms = self.positions[atoms_ind.ravel(),:] - abc_atoms[:,inds_tile[0]] += a_ind.ravel() - abc_atoms[:,inds_tile[1]] += b_ind.ravel() + abc_atoms = self.positions[atoms_ind.ravel(), :] + 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 = self.numbers[atoms_ind.ravel()] # Center atoms on image plane - 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 <= -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 = 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 <= -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) y = np.delete(y, atoms_del) atoms_ID_all = np.delete(atoms_ID_all, atoms_del) # Coordinate system for atomic projected potentials potential_radius = np.ceil(potential_radius_angstroms / pixel_size_angstroms) - R = np.arange(0.5-potential_radius,potential_radius+0.5) - R_ind = R.astype('int') - R_2D = np.sqrt(R[:,None]**2 + R[None,:]**2) + R = np.arange(0.5 - potential_radius, potential_radius + 0.5) + R_ind = R.astype("int") + R_2D = np.sqrt(R[:, None] ** 2 + R[None, :] ** 2) # Lookup table for atomic projected potentials atoms_ID = np.unique(self.numbers) - atoms_lookup = np.zeros(( - atoms_ID.shape[0], - R_2D.shape[0], - R_2D.shape[1], - )) + atoms_lookup = np.zeros( + ( + atoms_ID.shape[0], + R_2D.shape[0], + R_2D.shape[1], + ) + ) 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) + atoms_lookup[a0, :, :] = atom_sf.projected_potential(atoms_ID[a0], R_2D) atoms_lookup **= power_scale - # Thickness + # 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.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] + 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) @@ -1145,8 +1142,8 @@ def generate_projected_potential( 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_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], @@ -1156,11 +1153,13 @@ def generate_projected_potential( y_ind < im_size[1], ) - im_potential[x_ind[x_sub][:,None],y_ind[y_sub][None,:]] += atoms_lookup[ind][x_sub,:][:,y_sub] - + 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_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], @@ -1170,7 +1169,9 @@ def generate_projected_potential( y_ind < im_size[1], ) - im_potential[x_ind[x_sub][:,None],y_ind[y_sub][None,:]] += atoms_lookup[ind][x_sub,:][:,y_sub] + 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 @@ -1181,26 +1182,28 @@ def generate_projected_potential( im_potential = gaussian_filter( im_potential, sigma_image_blur, - mode = 'nearest', - ) + 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.98*int_vals.size).astype('int')], - )) + int_range = np.array( + ( + int_vals[np.round(0.02 * int_vals.size).astype("int")], + int_vals[np.round(0.98 * int_vals.size).astype("int")], + ) + ) - fig,ax = plt.subplots(figsize = figsize) + fig, ax = plt.subplots(figsize=figsize) ax.imshow( im_potential, - cmap = 'turbo', - vmin = int_range[0], - vmax = int_range[1], - ) + cmap="turbo", + vmin=int_range[0], + vmax=int_range[1], + ) ax.set_axis_off() - ax.set_aspect('equal') + ax.set_aspect("equal") return im_potential diff --git a/py4DSTEM/process/utils/single_atom_scatter.py b/py4DSTEM/process/utils/single_atom_scatter.py index 9085afd93..90397560f 100644 --- a/py4DSTEM/process/utils/single_atom_scatter.py +++ b/py4DSTEM/process/utils/single_atom_scatter.py @@ -2,6 +2,7 @@ import os from scipy.special import kn + class single_atom_scatter(object): """ This class calculates the composition averaged single atom scattering factor for a @@ -58,18 +59,17 @@ def projected_potential(self, Z, R): qe = 1.60217662e-19 # Permittivity of vacuum eps_0 = 8.85418782e-12 - # Bohr's constant + # Bohr's constant a_0 = 5.29177210903e-11 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)) + pre = 2 * np.pi / bi[i] ** 0.5 + fe += (ai[i] / bi[i] ** 1.5) * (kn(0, pre * R) + R * kn(1, pre * R)) # kappa = (4*np.pi*eps_0) / (2*np.pi*a_0*me) - return fe * 2 * np.pi**2# / kappa + return fe * 2 * np.pi**2 # / kappa # if units == "VA": # return h**2 / (2 * np.pi * me * qe) * 1e18 * fe # elif units == "A":