Adding support for beam precession

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,118 @@
     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(phi, r_spot, wavelength, theta):
+    """
+    Returns the z-coordinate in reciprocal space of the Ewald sphere at
+    distance r from the origin along the x-axis
+
+    Parameters
+    ----------
+    phi : float
+        The angle the beam is currently precessed to in degrees.
+        If the beam were projected on the x-y plane, the angle
+        is the angle between this projection with the x-axis.
+    r_spot : float
+        The distance to the point on the x-axis where we calculate z in A^-1
+    wavelength : float
+        The electron wavelength in A^-1
+    theta : 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
+    """
+    phi = np.deg2rad(phi)
+    r = 1/wavelength
+    theta = np.deg2rad(theta)
+    return (-np.sqrt(r**2*(1-np.sin(theta)**2*np.sin(phi)**2) -
+                     (r_spot - r*np.sin(theta)*np.cos(phi))**2) +
+            r*np.cos(theta))
+
+
+def _shape_factor_precession(z_spot, r_spot, wavelength, precession_angle,
+                             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
+    precession_angle : float
+        The precession angle in degrees
+    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
+        Maximum excitation angle to consider if precession_angle = 0. With
+        precession, it is rather a parameter to describe the "width" of the
+        rel-rods.
+
+    Other parameters
+    ----------------
+    ** kwargs: passed directly to 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
+    """
+    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, precession_angle)
+            return 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 +164,69 @@ 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.    
+    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.lorentzian`. 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=None,
+        shape_factor_model=None,
+        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
+        if self.precession_angle is None:
+            self.precession_angle = 0
+        if self.precession_angle < 0:
+            raise ValueError("The precession angle must be >= 0")
+        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."
+                )
+        elif callable(shape_factor_model):
+            self.shape_factor_model = shape_factor_model
+        elif shape_factor_model is None:
+            self.shape_factor_model = lorentzian
+        else:
+            raise TypeError(
+                "shape_factor_model must be a recognized string "
+                "or a callable object.")
+        self.minimum_intensity = minimum_intensity
+        self.shape_factor_kwargs = kwargs
         if scattering_params in ["lobato", "xtables"]:
             self.scattering_params = scattering_params
         else:
@@ -87,10 +241,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 +260,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 +297,50 @@ 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,
+                    self.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 +349,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]

diffsims/generators/library_generator.py

@@ -51,6 +51,8 @@ def get_diffraction_library(
         reciprocal_radius,
         half_shape,
         with_direct_beam=True,
+        max_excitation_error=1e-2,
+        debye_waller_factors={},
     ):
         """Calculates a dictionary of diffraction data for a library of crystal
         structures and orientations.
@@ -74,6 +76,11 @@ def get_diffraction_library(
             The half shape of the target patterns, for 144x144 use (72,72) etc
         with_direct_beam : bool
             Include the direct beam in the library.
+        max_excitation_error : float
+            The extinction distance for reflections, in reciprocal
+            Angstroms.
+        debye_waller_factors : dict of str:value pairs
+            Maps element names to their temperature-dependent Debye-Waller factors.
 
         Returns
         -------
@@ -103,6 +110,8 @@ def get_diffraction_library(
                     reciprocal_radius,
                     rotation=orientation,
                     with_direct_beam=with_direct_beam,
+                    max_excitation_error=max_excitation_error,
+                    debye_waller_factors=debye_waller_factors
                 )
 
                 # Calibrate simulation
@@ -117,12 +126,11 @@ def get_diffraction_library(
                 intensities[i] = simulation.intensities
 
             diffraction_library[phase_name] = {
-                "simulations": simulations,
-                "orientations": orientations,
-                "pixel_coords": pixel_coords,
-                "intensities": intensities,
-            }
-
+                 "simulations": simulations,
+                 "orientations": orientations,
+                 "pixel_coords": pixel_coords,
+                 "intensities": intensities,
+             }
         # Pass attributes to diffraction library from structure library.
         diffraction_library.identifiers = structure_library.identifiers
         diffraction_library.structures = structure_library.structures

diffsims/generators/sphere_mesh_generators.py

@@ -481,5 +481,5 @@ def beam_directions_grid_to_euler(vectors):
     phi2 = sign * np.nan_to_num(np.arccos(x_comp / norm_proj))
     # phi1 is just 0, rotation around z''
     phi1 = np.zeros(phi2.shape[0])
-    grid = np.rad2deg(np.vstack([phi1, Phi, phi2]).T)
+    grid = np.rad2deg(np.vstack([phi1, Phi, np.pi/2 - phi2]).T)
     return grid

diffsims/tests/test_generators/test_sphere_mesh_generators.py

@@ -86,5 +86,5 @@ def test_icosahedral_grid():
 
 def test_vectors_to_euler():
     grid = _normalize_vectors(np.array([[1, 0, 0], [0, 1, 0], [0, 1, 1], [1, 0, 1],]))
-    ang = np.array([[0, 90, 0], [0, 90, 90], [0, 45, 90], [0, 45, 0],])
+    ang = np.array([[0, 90, 90], [0, 90, 0], [0, 45, 0], [0, 45, 90],])
     assert np.allclose(ang, beam_directions_grid_to_euler(grid))