necessary functions in image_tools

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