HamishGBrown / py_multislice

A GPU accelerated Python multislice slice code
24 stars 16 forks source link

Precession Electron Diffraction simulations #7

Open anderscmathisen opened 1 year ago

anderscmathisen commented 1 year ago

Hi,

I am working on scanning precession electron diffraction (SPED, also known as double-conical beam-rocking) and have recently been using Pyms to simulate precession electron diffraction (PED) patterns of ErMnO3 and thought I would share my method in case anyone wants use it or further develop it. (I am still learning to use Pyms, and so feel free to correct me if I made any errors here)

Multislice simulations have been used for PED simulations in the past. The strategy I used with Pyms is the same as C.S. Own (https://onlinelibrary.wiley.com/iucr/doi/10.1107/S0108767306032892) used, which was to uniformly and discreetly sample many beam tilts from a cone with semi-angle of the precession angle and run the multislice simulation for each beam tilt, after which the results were incoherently summed to form the PED simulation. With Pyms, I used the CBED simulation in the Premixed_routines for this. With a beam tilt, the diffraction pattern shifts away from the center, and so the CBED pattern from each beam tilt simulation needs to be shifted back so that the direct beam is again in the center before the summation. Below is an example code showing this:

import pyms
import numpy as np
from pyms.Probe import wavev
import tqdm
import torch

#Load the crystal
crystal = pyms.structure.fromfile(some_file,
                    temperature_factor_units='B', 
                    atomic_coordinates = "fractional")

#Function to sample N beam tilts with radius in mrad
def beam_tilts(N, radius):
    thethas = np.linspace(0,360,N+1)

    xs = radius*np.cos(np.deg2rad(thethas))
    ys = radius*np.sin(np.deg2rad(thethas))

    arr =  np.array([xs,ys]).T

    return arr[:-1]

#Function for recentering simulated CBED. Sets pixels shifted outside to zero.
def shift_image(X, dx, dy):
    X = np.roll(X, dy, axis=1)
    X = np.roll(X, dx, axis=2)
    if dy>0:
        X[:,:dy, :] = 0
    elif dy<0:
        X[:,dy:, :] = 0
    if dx>0:
        X[:,:, :dx] = 0
    elif dx<0:
        X[:,:, dx:] = 0
    return X

# Set parameters for simulation
gridshape = [2048,2048]
tiling = [32,32]
subslices = [0.5,1]
eV = 200e3
app = 1.23
nfph = 25
device = torch.device('cpu')
thicknesses = np.linspace(250,1500,300)

beam_tilt_samples = 100
precession_angle = 17.453293 #in mrad,  = 1 deg precession

PED_simulation = np.zeros((300, 1365, 1365)) #array for storing results

#loop over beam tilts
for tilt in tqdm(beam_tilts(beam_tilt_samples, radius = precession_angle)):
    #do CBED simulation with beam tilt
    output = pyms.CBED(
            crystal,
            gridshape,
            eV,
            app,
            thicknesses,
            beam_tilt = tilt,
            tilt_units = "mrad",
            subslices=subslices,
            tiling=tiling,
            nfph=nfph,
            showProgress=False,
            device_type=device,
        )

    output = np.array(output)

    #calculate pixel shift to recenter CBED pattern
    k = wavev(eV)
    tilt_ = np.asarray(tilt) * 1e-3 * k
    shift = tilt_ *crystal.unitcell[:2] * np.asarray(tiling)
    x, y = shift

    #recenter CBED simulation  
    output = shift_image(output,int(-y),int(-x))

    #add simulation to array
    PED_simulation += output

#Code to save or analyze results

From my testing, this seems to work well for simulating PED, with results matching decently with my experimental data. It is however quite slow as computation time increases linearly with number of beam tilts, and so improvements in terms of computational efficiency would be nice (Yes, I know cuda would make it much faster, but I dont have one available :().

HamishGBrown commented 1 year ago

Hi Anders, thanks for using py_multislice and sorry it has taken me so long to get back to this message. Unfortunately there's not much to be done in terms of beating the linear-scaling nature of this calculation since every different beam tilt for a object with dynamical diffraction (ie. one you would use precession on) produces a different result. One thing that might help is to make the multislice transmission functions (these are the exponentials of the electrostatic potentials of your sample) and propagators ahead of time using the helper function multislice_precursor and pass them to the CBED function. The default is to not do this and in this case the CBED function will generate them internally each time you call it. Without testing something like this should work if you add the following outside your loop over tilts:

propagators,transmissionfns = pyms.multislice_precursor(crystal,gridshape,eV,subslices=[0.5,0.1],tiling=tiling)

Then when you call CBED:

output = pyms.CBED(
            crystal,
            gridshape,
            eV,
            app,
            thicknesses,
            beam_tilt = tilt,
            tilt_units = "mrad",
            subslices=[0.5,1.0],
            tiling=tiling,
            nfph=nfph,
            showProgress=False,
            device_type=device,
            P=propagators,
            T=transmissionfns,
        )

Note that I've added "sub-slicing" to the multislice calculation in the above, basically if your unit cell is larger than 2 Angstrom then you should slice it depth-ways