From ce48099eb3501a9a430adc0c410beec3af4af7ce Mon Sep 17 00:00:00 2001 From: ahoust17 <88668350+ahoust17@users.noreply.github.com> Date: Sat, 14 Dec 2024 14:39:58 -0500 Subject: [PATCH] necessary functions in image_tools --- pyTEMlib/image_tools.py | 110 ++++++++++++++++++++++------------------ 1 file changed, 62 insertions(+), 48 deletions(-) diff --git a/pyTEMlib/image_tools.py b/pyTEMlib/image_tools.py index 028d6b04..76325e85 100644 --- a/pyTEMlib/image_tools.py +++ b/pyTEMlib/image_tools.py @@ -56,7 +56,11 @@ from sklearn.cluster import DBSCAN from ase.build import fcc110 -import probe_tools # Assuming you have this module available +from pyTEMlib import probe_tools + +from scipy.ndimage import rotate +from scipy.interpolate import RegularGridInterpolator +from scipy.signal import fftconvolve _SimpleITK_present = True @@ -71,60 +75,70 @@ 'install with: conda install -c simpleitk simpleitk ') +def get_atomic_pseudo_potential(fov, atoms, size=512, rotation=0): + # Big assumption: the atoms are not near the edge of the unit cell + # If any atoms are close to the edge (ex. [0,0]) then the potential will be clipped + # before calling the function, shift the atoms to the center of the unit cell + + pixel_size = fov / size + max_size = int(size * np.sqrt(2) + 1) # Maximum size to accommodate rotation + + # Create unit cell potential + positions = atoms.get_positions()[:, :2] + atomic_numbers = atoms.get_atomic_numbers() + unit_cell_size = atoms.cell.cellpar()[:2] + + unit_cell_potential = np.zeros((max_size, max_size)) + for pos, atomic_number in zip(positions, atomic_numbers): + x = pos[0] / pixel_size + y = pos[1] / pixel_size + atom_width = 0.5 # Angstrom + gauss_width = atom_width/pixel_size # important for images at various fov. Room for improvement with theory + gauss = probe_tools.make_gauss(max_size, max_size, width = gauss_width, x0=x, y0=y) + unit_cell_potential += gauss * atomic_number # gauss is already normalized to 1 + + # Create interpolation function for unit cell potential + x_grid = np.linspace(0, fov * max_size / size, max_size) + y_grid = np.linspace(0, fov * max_size / size, max_size) + interpolator = RegularGridInterpolator((x_grid, y_grid), unit_cell_potential, bounds_error=False, fill_value=0) + + # Vectorized computation of the full potential map with max_size + x_coords, y_coords = np.meshgrid(np.linspace(0, fov, max_size), np.linspace(0, fov, max_size), indexing="ij") + xtal_x = x_coords % unit_cell_size[0] + xtal_y = y_coords % unit_cell_size[1] + potential_map = interpolator((xtal_x.ravel(), xtal_y.ravel())).reshape(max_size, max_size) + + # Rotate and crop the potential map + potential_map = rotate(potential_map, rotation, reshape=False) + center = potential_map.shape[0] // 2 + potential_map = potential_map[center - size // 2:center + size // 2, center - size // 2:center + size // 2] -def simulate_atomic_scattering_potential(fov=100, angle=0, atoms=None): - """ - Simulates the atomic scattering potential on a 2D grid. + potential_map = scipy.ndimage.gaussian_filter(potential_map,3) - Parameters: - - fov (float): Field of view in angstroms. - - angle (float): Rotation angle in degrees. - - atoms (ase.Atoms): Optional input for a custom atomic structure. Defaults to Al FCC(110). + return potential_map - Returns: - - np.ndarray: 2D grid of the simulated scattering potential. - """ - # Default to Al FCC(110) structure if no atoms are provided - if atoms is None: - atoms = fcc110('Al', size=(2, 2, 1), orthogonal=True) - - # Ensure the structure is periodic and centered - atoms.pbc = True - atoms.center() +def convolve_probe(ab, potential): + # the pixel sizes should be the exact same as the potential + final_sizes = potential.shape + + # Perform FFT-based convolution + pad_height = pad_width = potential.shape[0] // 2 + potential = np.pad(potential, ((pad_height, pad_height), (pad_width, pad_width)), mode='constant') + + probe, A_k, chi = probe_tools.get_probe(ab, potential.shape[0], potential.shape[1], scale = 'mrad', verbose= False) - # Determine the grid size based on the field of view - grid_scaling_factor = int(1 / (fov / 1024)) - grid_x = int(atoms.cell[0, 0] * grid_scaling_factor) - grid_y = int(atoms.cell[1, 1] * grid_scaling_factor) - - # Create an empty grid and calculate pixel size - potential_map = np.zeros((grid_x, grid_y)) - pixel_size_x = atoms.cell[0, 0] / grid_x - pixel_size_y = atoms.cell[1, 1] / grid_y - - # Generate the scattering potential by summing Gaussian peaks - for position in atoms.positions: - x_pixel = position[0] / pixel_size_x - y_pixel = position[1] / pixel_size_y - gauss_peak = probe_tools.make_gauss(grid_x, grid_y, x0=x_pixel, y0=y_pixel) - potential_map += gauss_peak - - # Expand the grid to ensure proper rotation without clipping - expansion_factor = np.floor(1450 / np.array(potential_map.shape)).astype(int) - potential_map = np.tile(potential_map, (expansion_factor[0], expansion_factor[1])) - - # Rotate the map by the specified angle - potential_map = scipy.ndimage.rotate(potential_map, angle, axes=(1, 0), reshape=True) - - # Extract a centered square region of size 1024x1024 - center = potential_map.shape[0] // 2 - cropped_map = potential_map[center - 512:center + 512, center - 512:center + 512] - # Debugging output for pixel size - print(f"Pixel size (angstrom): x={pixel_size_x}, y={pixel_size_y}") + convolved = fftconvolve(potential, probe, mode='same') + + # Crop to original potential size + start_row = pad_height + start_col = pad_width + end_row = start_row + final_sizes[0] + end_col = start_col + final_sizes[1] - return cropped_map + image = convolved[start_row:end_row, start_col:end_col] + return probe, image # Wavelength in 1/nm