Merge pull request #137 from din14970/update_dp_sim

Adding support for beam precession

CHANGELOG.md

@@ -6,11 +6,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Changed
-- get_grid_beam_directions, now works based off of meshes
+- `get_grid_beam_directions`, now works based off of meshes
+- the arguments in the `DiffractionGenerator` constructor and the `DiffractionLibraryGenerator.get_diffraction_library` function have been shuffled so that the former captures arguments related to "the instrument/physics" while the latter captures arguments relevant to "the sample/material".
 
 ### Added
 - API reference documentation via Read The Docs: https://diffsims.rtfd.io
-- New module: "sphere_mesh_generators"
+- New module: `sphere_mesh_generators`
+- beam precession is now supported in simulating electron diffraction patterns
+- more shape factor functions have been added
 - This project now keeps a Changelog
 
 ### Removed

diffsims/generators/diffraction_generator.py

@@ -19,6 +19,7 @@
 """Electron diffraction pattern simulation."""
 
 import numpy as np
+from scipy.integrate import quad
 from transforms3d.euler import euler2mat
 
 from diffsims.sims.diffraction_simulation import DiffractionSimulation
@@ -31,7 +32,130 @@
     get_intensities_params,
 )
 from diffsims.utils.fourier_transform import from_recip
-from diffsims.utils.shape_factor_models import linear
+from diffsims.utils.shape_factor_models import (
+    linear,
+    atanc,
+    lorentzian,
+    sinc,
+    sin2c,
+    lorentzian_precession,
+)
+
+
+_shape_factor_model_mapping = {
+    "linear": linear,
+    "atanc": atanc,
+    "sinc": sinc,
+    "sin2c": sin2c,
+    "lorentzian": lorentzian,
+}
+
+
+def _z_sphere_precession(theta, r_spot, wavelength, phi):
+    """
+    Returns the z-coordinate in reciprocal space of the Ewald sphere at
+    distance r_spot from the origin
+
+    Parameters
+    ----------
+    theta : float
+        The azimuthal angle in degrees
+    r_spot : float
+        The projected length of the reciprocal lattice vector onto the plane
+        perpendicular to the optical axis in A^-1
+    wavelength : float
+        The electron wavelength in A
+    phi : float
+        The precession angle in degrees (angle between beam and optical axis)
+
+    Returns
+    -------
+    z : float
+        The height of the ewald sphere at the point r in A^-1
+
+    Notes
+    -----
+    * The azimuthal angle is the angle the beam is currently precessed to.
+    It is the angle between the projection of the beam and the projection of
+    the relevant diffraction spot both onto the x-y plane.
+    * In the derivation of this formula we assume that we will always integrate
+    over a full precession circle, because we do not explicitly take into
+    consideration x-y coordinates of reflections.
+    """
+    theta = np.deg2rad(theta)
+    r = 1 / wavelength
+    phi = np.deg2rad(phi)
+    return -np.sqrt(
+        r ** 2 * (1 - np.sin(phi) ** 2 * np.sin(theta) ** 2)
+        - (r_spot - r * np.sin(phi) * np.cos(theta)) ** 2
+    ) + r * np.cos(phi)
+
+
+def _shape_factor_precession(
+    z_spot, r_spot, wavelength, phi, shape_function, max_excitation, **kwargs
+):
+    """
+    The rel-rod shape factors for reflections taking into account
+    precession
+
+    Parameters
+    ----------
+    z_spot : np.ndarray (N,)
+        An array representing the z-coordinates of the reflections in A^-1
+    r_spot : np.ndarray (N,)
+        An array representing the distance of spots from the z-axis in A^-1
+    wavelength : float
+        The electron wavelength in A
+    phi : float
+        The precession angle in degrees
+    shape_function : callable
+        A function that describes the influence from the rel-rods. Should be
+        in the form func(excitation_error: np.ndarray, max_excitation: float,
+        **kwargs)
+    max_excitation : float
+        Parameter to describe the "extent" of the rel-rods.
+
+    Other parameters
+    ----------------
+    ** kwargs: passed directly to shape_function
+
+    Notes
+    -----
+    * We calculate excitation_error as z_spot - z_sphere so that it is
+    negative when the spot is outside the ewald sphere and positive when inside
+    conform W&C chapter 12, section 12.6
+    * We assume that the sample is a thin infinitely wide slab perpendicular
+    to the optical axis, so that the shape factor function only depends on the
+    distance from each spot to the Ewald sphere parallel to the optical axis.
+    """
+    shf = np.zeros(z_spot.shape)
+    # loop over all spots
+    for i, (z_spot_i, r_spot_i) in enumerate(zip(z_spot, r_spot)):
+
+        def integrand(phi):
+            z_sph = _z_sphere_precession(phi, r_spot_i, wavelength, phi)
+            return shape_function(z_spot_i - z_sph, max_excitation, **kwargs)
+
+        # average factor integrated over the full revolution of the beam
+        shf[i] = (1 / (360)) * quad(integrand, 0, 360)[0]
+    return shf
+
+
+def _average_excitation_error_precession(z_spot, r_spot, wavelength, precession_angle):
+    """
+    Calculate the average excitation error for spots
+    """
+    ext = np.zeros(z_spot.shape)
+    # loop over all spots
+    for i, (z_spot_i, r_spot_i) in enumerate(zip(z_spot, r_spot)):
+
+        def integrand(phi):
+            z_sph = _z_sphere_precession(phi, r_spot_i, wavelength, precession_angle)
+            return z_spot_i - z_sph
+
+        # average factor integrated over the full revolution of the beam
+        ext[i] = (1 / (360)) * quad(integrand, 0, 360)[0]
+    return ext
 
 
 class DiffractionGenerator(object):
@@ -52,27 +176,61 @@ class DiffractionGenerator(object):
     ----------
     accelerating_voltage : float
         The accelerating voltage of the microscope in kV.
-    max_excitation_error : float
-        Removed in this version, defaults to None
-    debye_waller_factors : dict of str:value pairs
-        Maps element names to their temperature-dependent Debye-Waller factors.
     scattering_params : str
         "lobato" or "xtables"
+    minimum_intensity : float
+        Minimum intensity for a peak to be considered visible in the pattern
+    precession_angle : float
+        Angle about which the beam is precessed. Default is no precession.
+    approximate_precession : boolean
+        When using precession, whether to precisely calculate average
+        excitation errors and intensities or use an approximation. See notes.
+    shape_factor_model : function or string
+        A function that takes excitation_error and
+        `max_excitation_error` (and potentially kwargs) and returns
+        an intensity scaling factor. If None defaults to
+        `shape_factor_models.linear`. A number of pre-programmed functions
+        are available via strings.
+    kwargs
+        Keyword arguments passed to `shape_factor_model`.
+
+    Notes
+    -----
+    * A full calculation is much slower and is not recommended for calculating
+      a diffraction library for precession diffraction patterns.
+    * When using precession and approximate_precession=True, the shape factor
+    model defaults to Lorentzian; shape_factor_model is ignored. Only with
+    approximate_precession=False the custom shape_factor_model is used.
     """
 
     def __init__(
         self,
         accelerating_voltage,
-        max_excitation_error=None,
-        debye_waller_factors={},
         scattering_params="lobato",
+        precession_angle=0,
+        shape_factor_model="lorentzian",
+        approximate_precession=True,
+        minimum_intensity=1e-20,
+        **kwargs,
     ):
-        if max_excitation_error is not None:
-            print(
-                "This class changed in v0.3 and no longer takes a maximum_excitation_error"
-            )
         self.wavelength = get_electron_wavelength(accelerating_voltage)
-        self.debye_waller_factors = debye_waller_factors
+        self.precession_angle = precession_angle
+        self.approximate_precession = approximate_precession
+        if isinstance(shape_factor_model, str):
+            if shape_factor_model in _shape_factor_model_mapping.keys():
+                self.shape_factor_model = _shape_factor_model_mapping[
+                    shape_factor_model
+                ]
+            else:
+                raise NotImplementedError(
+                    f"{shape_factor_model} is not a recognized shape factor "
+                    f"model, choose from: {_shape_factor_model_mapping.keys()} "
+                    f"or provide your own function."
+                )
+        else:
+            self.shape_factor_model = shape_factor_model
+        self.minimum_intensity = minimum_intensity
+        self.shape_factor_kwargs = kwargs
         if scattering_params in ["lobato", "xtables"]:
             self.scattering_params = scattering_params
         else:
@@ -87,10 +245,9 @@ def calculate_ed_data(
         structure,
         reciprocal_radius,
         rotation=(0, 0, 0),
-        shape_factor_model=None,
-        max_excitation_error=1e-2,
         with_direct_beam=True,
-        **kwargs
+        max_excitation_error=1e-2,
+        debye_waller_factors={},
     ):
         """Calculates the Electron Diffraction data for a structure.
 
@@ -107,19 +264,14 @@ def calculate_ed_data(
         rotation : tuple
             Euler angles, in degrees, in the rzxz convention. Default is
             (0, 0, 0) which aligns 'z' with the electron beam.
-        shape_factor_model : function or None
-            A function that takes excitation_error and
-            `max_excitation_error` (and potentially kwargs) and returns
-            an intensity scaling factor. If None defaults to
-            `shape_factor_models.linear`.
-        max_excitation_error : float
-            The extinction distance for reflections, in reciprocal
-            Angstroms.
         with_direct_beam : bool
             If True, the direct beam is included in the simulated
             diffraction pattern. If False, it is not.
-        kwargs
-            Keyword arguments passed to `shape_factor_model`.
+        max_excitation_error : float
+            The extinction distance for reflections, in reciprocal
+            Angstroms. Roughly equal to 1/thickness.
+        debye_waller_factors : dict of str:value pairs
+            Maps element names to their temperature-dependent Debye-Waller factors.
 
         Returns
         -------
@@ -149,21 +301,49 @@ def calculate_ed_data(
         # excitation error and store the magnitude of their excitation error.
         r_sphere = 1 / wavelength
         r_spot = np.sqrt(np.sum(np.square(cartesian_coordinates[:, :2]), axis=1))
-        z_sphere = -np.sqrt(r_sphere ** 2 - r_spot ** 2) + r_sphere
-        excitation_error = np.absolute(z_sphere - cartesian_coordinates[:, 2])
-        intersection = excitation_error < max_excitation_error
+        z_spot = cartesian_coordinates[:, 2]
+
+        if self.precession_angle > 0 and not self.approximate_precession:
+            # We find the average excitation error - this step can be
+            # quite expensive
+            excitation_error = _average_excitation_error_precession(
+                z_spot,
+                r_spot,
+                wavelength,
+                self.precession_angle,
+            )
+        else:
+            z_sphere = -np.sqrt(r_sphere ** 2 - r_spot ** 2) + r_sphere
+            excitation_error = z_sphere - z_spot
         # Mask parameters corresponding to excited reflections.
+        intersection = np.abs(excitation_error) < max_excitation_error
         intersection_coordinates = cartesian_coordinates[intersection]
-        g_indices = spot_indices[intersection]
         excitation_error = excitation_error[intersection]
+        r_spot = r_spot[intersection]
+        g_indices = spot_indices[intersection]
         g_hkls = spot_distances[intersection]
-
-        if shape_factor_model is not None:
-            shape_factor = shape_factor_model(
-                excitation_error, max_excitation_error, **kwargs
+        # take into consideration rel-rods
+        if self.precession_angle > 0 and not self.approximate_precession:
+            shape_factor = _shape_factor_precession(
+                intersection_coordinates[:, 2],
+                r_spot,
+                wavelength,
+                self.precession_angle,
+                self.shape_factor_model,
+                max_excitation_error,
+                **self.shape_factor_kwargs,
+            )
+        elif self.precession_angle > 0 and self.approximate_precession:
+            shape_factor = lorentzian_precession(
+                excitation_error,
+                max_excitation_error,
+                r_spot,
+                self.precession_angle,
             )
         else:
-            shape_factor = linear(excitation_error, max_excitation_error)
+            shape_factor = self.shape_factor_model(
+                excitation_error, max_excitation_error, **self.shape_factor_kwargs
+            )
 
         # Calculate diffracted intensities based on a kinematical model.
         intensities = get_kinematical_intensities(
@@ -172,11 +352,11 @@ def calculate_ed_data(
             g_hkls,
             prefactor=shape_factor,
             scattering_params=self.scattering_params,
-            debye_waller_factors=self.debye_waller_factors,
+            debye_waller_factors=debye_waller_factors,
         )
 
         # Threshold peaks included in simulation based on minimum intensity.
-        peak_mask = intensities > 1e-20
+        peak_mask = intensities > self.minimum_intensity
         intensities = intensities[peak_mask]
         intersection_coordinates = intersection_coordinates[peak_mask]
         g_indices = g_indices[peak_mask]
@@ -189,7 +369,11 @@ def calculate_ed_data(
         )
 
     def calculate_profile_data(
-        self, structure, reciprocal_radius=1.0, minimum_intensity=1e-3
+        self,
+        structure,
+        reciprocal_radius=1.0,
+        minimum_intensity=1e-3,
+        debye_waller_factors={},
     ):
         """Calculates a one dimensional diffraction profile for a
         structure.
@@ -204,6 +388,8 @@ def calculate_profile_data(
         minimum_intensity : float
             The minimum intensity required for a diffraction peak to be
             considered real. Deals with numerical precision issues.
+        debye_waller_factors : dict of str:value pairs
+            Maps element names to their temperature-dependent Debye-Waller factors.
 
         Returns
         -------
@@ -231,7 +417,7 @@ def calculate_profile_data(
             np.asarray(g_hkls),
             prefactor=multiplicities,
             scattering_params=self.scattering_params,
-            debye_waller_factors=self.debye_waller_factors,
+            debye_waller_factors=debye_waller_factors,
         )
 
         if is_lattice_hexagonal(latt):
@@ -299,7 +485,7 @@ def calculate_ed_data(
         dtype="float64",
         ZERO=1e-14,
         mode="kinematic",
-        **kwargs
+        **kwargs,
     ):
         """
         Calculates single electron diffraction image for particular atomic