From 876aed2ab731be65b657d22e38f9a54362e16ef3 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Fri, 13 Nov 2020 13:52:41 +0100
Subject: [PATCH 01/10] fix small bug orientations

---
 diffsims/generators/sphere_mesh_generators.py                 | 2 +-
 diffsims/tests/test_generators/test_sphere_mesh_generators.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/diffsims/generators/sphere_mesh_generators.py b/diffsims/generators/sphere_mesh_generators.py
index 2159f5c5..1e2b289a 100644
--- a/diffsims/generators/sphere_mesh_generators.py
+++ b/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
diff --git a/diffsims/tests/test_generators/test_sphere_mesh_generators.py b/diffsims/tests/test_generators/test_sphere_mesh_generators.py
index 06d02b05..1956a8c1 100644
--- a/diffsims/tests/test_generators/test_sphere_mesh_generators.py
+++ b/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))

From 9735db80fdc5b4efa0e53ee6f6437d93e40bd7af Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Sun, 15 Nov 2020 12:02:05 +0100
Subject: [PATCH 02/10] added support for precession

---
 diffsims/generators/diffraction_generator.py | 152 +++++++++++++++++--
 diffsims/utils/shape_factor_models.py        |  55 ++++++-
 2 files changed, 195 insertions(+), 12 deletions(-)

diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index 76e7a9bb..ae2185bc 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/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
@@ -34,6 +35,101 @@
 from diffsims.utils.shape_factor_models import linear
 
 
+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):
     """Computes electron diffraction patterns for a crystal structure.
 
@@ -90,6 +186,9 @@ def calculate_ed_data(
         shape_factor_model=None,
         max_excitation_error=1e-2,
         with_direct_beam=True,
+        minimum_intensity=1e-20,
+        precession_angle=0,
+        full_calculation=False,
         **kwargs
     ):
         """Calculates the Electron Diffraction data for a structure.
@@ -111,13 +210,20 @@ def calculate_ed_data(
             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`.
+            `shape_factor_models.atanc`.
         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.
+        minimum_intensity : float
+            Minimum intensity for a peak to be considered in the pattern
+        precession_angle : float
+            Angle about which the beam is precessed
+        full_calculation : boolean
+            When using precession, whether to precisely calculate average
+            excitation errors for filtering diffraction spots.
         kwargs
             Keyword arguments passed to `shape_factor_model`.
 
@@ -149,21 +255,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 precession_angle is None:
+            precession_angle = 0
+
+        if shape_factor_model is None:
+            shape_factor_model = linear
+
+        if precession_angle > 0 and full_calculation:
+            # We find the average excitation error - this step can be
+            # quite expensive
+            excitation_error = _average_excitation_error_precession(
+                    z_spot,
+                    r_spot,
+                    wavelength,
+                    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]
+        g_indices = spot_indices[intersection]
         g_hkls = spot_distances[intersection]
-
-        if shape_factor_model is not None:
+        # take into consideration rel-rods
+        if precession_angle > 0:
+            shape_factor = _shape_factor_precession(
+                    intersection_coordinates[:, 2],
+                    r_spot[intersection],
+                    wavelength,
+                    precession_angle,
+                    shape_factor_model,
+                    max_excitation_error,
+                    **kwargs,
+                    )
+        elif precession_angle == 0:
             shape_factor = shape_factor_model(
                 excitation_error, max_excitation_error, **kwargs
             )
         else:
-            shape_factor = linear(excitation_error, max_excitation_error)
+            raise ValueError("The precession angle must be >= 0")
 
         # Calculate diffracted intensities based on a kinematical model.
         intensities = get_kinematical_intensities(
@@ -176,7 +310,7 @@ def calculate_ed_data(
         )
 
         # Threshold peaks included in simulation based on minimum intensity.
-        peak_mask = intensities > 1e-20
+        peak_mask = intensities > minimum_intensity
         intensities = intensities[peak_mask]
         intersection_coordinates = intersection_coordinates[peak_mask]
         g_indices = g_indices[peak_mask]
diff --git a/diffsims/utils/shape_factor_models.py b/diffsims/utils/shape_factor_models.py
index e7af7d65..0f103720 100644
--- a/diffsims/utils/shape_factor_models.py
+++ b/diffsims/utils/shape_factor_models.py
@@ -77,7 +77,56 @@ def sinc(excitation_error, max_excitation_error, minima_number=5):
     -------
     intensity : array-like or float
     """
+    fac = np.pi * minima_number / max_excitation_error
+    num = np.sin(fac * excitation_error)
+    denom = fac * excitation_error
+    return np.nan_to_num(
+            np.abs(np.divide(num, denom, out=np.zeros_like(num), where=denom != 0)),
+            nan=1,
+            )
 
-    num = np.sin(np.pi * minima_number * excitation_error / max_excitation_error)
-    denom = excitation_error
-    return np.abs(np.divide(num, denom, out=np.zeros_like(num), where=denom != 0))
+
+def sin2c(excitation_error, max_excitation_error, minima_number=5):
+    """
+    Intensity with sin^2(s)/s^2 profile, after Howie-Whelan rel-rod
+
+    Parameters
+    ----------
+    excitation_error : array-like or float
+        The distance (reciprocal) from a reflection to the Ewald sphere
+
+    max_excitation_error : float
+        The distance at which a reflection becomes extinct
+
+    minima_number : int
+        The minima_number'th minima lies at max_excitation_error from 0
+
+    Returns
+    -------
+    intensity : array-like or float
+    """
+    return sinc(excitation_error, max_excitation_error, minima_number)**2
+
+
+def atanc(excitation_error, max_excitation_error):
+    """
+    Intensity with arctan(s)/s profile that closely follows sin(s)/s but
+    is smooth for s!=0.
+
+    Parameters
+    ----------
+    excitation_error : array-like or float
+        The distance (reciprocal) from a reflection to the Ewald sphere
+
+    max_excitation_error : float
+        The distance at which a reflection becomes extinct
+
+    Returns
+    -------
+    intensity : array-like or float
+    """
+    fac = np.pi * 5 / np.abs(max_excitation_error)
+    return np.nan_to_num(
+            np.arctan(fac*excitation_error)/(fac*excitation_error),
+            nan=1,
+            )

From c957383ae3718557be473453f9866a1a6f474bdb Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Tue, 17 Nov 2020 15:52:48 +0100
Subject: [PATCH 03/10] shuffled parameters, implemented simplified lorentzian

---
 diffsims/generators/diffraction_generator.py | 155 ++++++++++++-------
 diffsims/generators/library_generator.py     |  20 ++-
 diffsims/utils/shape_factor_models.py        |  68 +++++++-
 3 files changed, 180 insertions(+), 63 deletions(-)

diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index ae2185bc..d80ea908 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/diffsims/generators/diffraction_generator.py
@@ -32,7 +32,23 @@
     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):
@@ -148,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:
@@ -183,13 +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,
-        minimum_intensity=1e-20,
-        precession_angle=0,
-        full_calculation=False,
-        **kwargs
+        max_excitation_error=1e-2,
+        debye_waller_factors={},
     ):
         """Calculates the Electron Diffraction data for a structure.
 
@@ -206,26 +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.atanc`.
-        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.
-        minimum_intensity : float
-            Minimum intensity for a peak to be considered in the pattern
-        precession_angle : float
-            Angle about which the beam is precessed
-        full_calculation : boolean
-            When using precession, whether to precisely calculate average
-            excitation errors for filtering diffraction spots.
-        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
         -------
@@ -257,20 +299,14 @@ def calculate_ed_data(
         r_spot = np.sqrt(np.sum(np.square(cartesian_coordinates[:, :2]), axis=1))
         z_spot = cartesian_coordinates[:, 2]
 
-        if precession_angle is None:
-            precession_angle = 0
-
-        if shape_factor_model is None:
-            shape_factor_model = linear
-
-        if precession_angle > 0 and full_calculation:
+        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,
-                    precession_angle,
+                    self.precession_angle,
                     )
         else:
             z_sphere = -np.sqrt(r_sphere ** 2 - r_spot ** 2) + r_sphere
@@ -279,25 +315,32 @@ def calculate_ed_data(
         intersection = np.abs(excitation_error) < max_excitation_error
         intersection_coordinates = cartesian_coordinates[intersection]
         excitation_error = excitation_error[intersection]
+        r_spot = r_spot[intersection]
         g_indices = spot_indices[intersection]
         g_hkls = spot_distances[intersection]
         # take into consideration rel-rods
-        if precession_angle > 0:
+        if self.precession_angle > 0 and not self.approximate_precession:
             shape_factor = _shape_factor_precession(
                     intersection_coordinates[:, 2],
-                    r_spot[intersection],
+                    r_spot,
                     wavelength,
-                    precession_angle,
-                    shape_factor_model,
-                    max_excitation_error,
-                    **kwargs,
+                    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,
                     )
-        elif precession_angle == 0:
-            shape_factor = shape_factor_model(
-                excitation_error, max_excitation_error, **kwargs
-            )
         else:
-            raise ValueError("The precession angle must be >= 0")
+            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(
@@ -306,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 > minimum_intensity
+        peak_mask = intensities > self.minimum_intensity
         intensities = intensities[peak_mask]
         intersection_coordinates = intersection_coordinates[peak_mask]
         g_indices = g_indices[peak_mask]
diff --git a/diffsims/generators/library_generator.py b/diffsims/generators/library_generator.py
index bbc0100e..e45b29ee 100644
--- a/diffsims/generators/library_generator.py
+++ b/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
diff --git a/diffsims/utils/shape_factor_models.py b/diffsims/utils/shape_factor_models.py
index 0f103720..6cbb30e3 100644
--- a/diffsims/utils/shape_factor_models.py
+++ b/diffsims/utils/shape_factor_models.py
@@ -55,7 +55,7 @@ def linear(excitation_error, max_excitation_error):
     intensities : array-like or float
     """
 
-    return 1 - excitation_error / max_excitation_error
+    return 1 - np.abs(excitation_error) / max_excitation_error
 
 
 def sinc(excitation_error, max_excitation_error, minima_number=5):
@@ -130,3 +130,69 @@ def atanc(excitation_error, max_excitation_error):
             np.arctan(fac*excitation_error)/(fac*excitation_error),
             nan=1,
             )
+
+
+def lorentzian(excitation_error, max_excitation_error):
+    """
+    Lorentzian intensity profile that should approximate
+    the two-beam rocking curve. After [1].
+
+    Parameters
+    ----------
+    excitation_error : array-like or float
+        The distance (reciprocal) from a reflection to the Ewald sphere
+
+    max_excitation_error : float
+        The distance at which a reflection becomes extinct
+
+    Returns
+    -------
+    intensity_factor : array-like or float
+        Vector representing the rel-rod factor for each reflection
+
+    Notes
+    -----
+    [1] L. Palatinus, P. Brázda, M. Jelínek, J. Hrdá, G. Steciuk, M. Klementová, Specifics of the data processing of precession electron diffraction tomography data and their implementation in the program PETS2.0, Acta Crystallogr. Sect. B Struct. Sci. Cryst. Eng. Mater. 75 (2019) 512–522. doi:10.1107/S2052520619007534.
+    """
+    # in the paper, sigma = pi*thickness.
+    # We assume thickness = 1/max_exitation_error
+    sigma = np.pi/max_excitation_error
+    fac = sigma/(np.pi*(sigma**2*excitation_error**2+1))
+    return fac
+
+
+def lorentzian_precession(excitation_error, max_excitation_error,
+                          r_spot, precession_angle):
+    """
+    Intensity profile factor for a precessed beam assuming a Lorentzian
+    intensity profile for the un-precessed beam. After [1].
+
+    Parameters
+    ----------
+    excitation_error : array-like or float
+        The distance (reciprocal) from a reflection to the Ewald sphere
+
+    max_excitation_error : float
+        The distance at which a reflection becomes extinct
+
+    r_spot : array-like or float
+        The distance (reciprocal) from each reflection to the origin
+
+    precession_angle : float
+        The beam precession angle in degrees; the angle the beam makes
+        with the optical axis.
+
+    Returns
+    -------
+    intensity_factor : array-like or float
+        Vector representing the rel-rod factor for each reflection
+
+    Notes
+    -----
+    [1] L. Palatinus, P. Brázda, M. Jelínek, J. Hrdá, G. Steciuk, M. Klementová, Specifics of the data processing of precession electron diffraction tomography data and their implementation in the program PETS2.0, Acta Crystallogr. Sect. B Struct. Sci. Cryst. Eng. Mater. 75 (2019) 512–522. doi:10.1107/S2052520619007534.
+    """
+    sigma = np.pi/max_excitation_error
+    u = sigma**2*(r_spot**2*precession_angle**2 - excitation_error**2)+1
+    z = np.sqrt(u**2 + 4*sigma**2 + excitation_error**2)
+    fac = (sigma/np.pi)*np.sqrt(2*(u+z)/z**2)
+    return fac

From 9e37d33a9946d96ec749b62810925a049423501d Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Tue, 17 Nov 2020 21:32:43 +0100
Subject: [PATCH 04/10] implemented review from pc494

---
 diffsims/generators/diffraction_generator.py | 37 ++++++++------------
 diffsims/utils/shape_factor_models.py        | 16 +++++----
 2 files changed, 24 insertions(+), 29 deletions(-)

diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index d80ea908..13ed8502 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/diffsims/generators/diffraction_generator.py
@@ -177,7 +177,7 @@ class DiffractionGenerator(object):
         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
+        `shape_factor_models.linear`. A number of pre-programmed functions
         are available via strings.
     kwargs
         Keyword arguments passed to `shape_factor_model`.
@@ -195,18 +195,14 @@ def __init__(
         self,
         accelerating_voltage,
         scattering_params="lobato",
-        precession_angle=None,
-        shape_factor_model=None,
+        precession_angle=0,
+        shape_factor_model="linear",
         approximate_precession=True,
         minimum_intensity=1e-20,
         **kwargs,
     ):
         self.wavelength = get_electron_wavelength(accelerating_voltage)
         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():
@@ -217,14 +213,8 @@ def __init__(
                     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.shape_factor_model = shape_factor_model
         self.minimum_intensity = minimum_intensity
         self.shape_factor_kwargs = kwargs
         if scattering_params in ["lobato", "xtables"]:
@@ -299,15 +289,16 @@ def calculate_ed_data(
         r_spot = np.sqrt(np.sum(np.square(cartesian_coordinates[:, :2]), axis=1))
         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,
-                    )
+        if self.precession_angle > 0:
+            if 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
diff --git a/diffsims/utils/shape_factor_models.py b/diffsims/utils/shape_factor_models.py
index 6cbb30e3..c14f4a00 100644
--- a/diffsims/utils/shape_factor_models.py
+++ b/diffsims/utils/shape_factor_models.py
@@ -108,7 +108,7 @@ def sin2c(excitation_error, max_excitation_error, minima_number=5):
     return sinc(excitation_error, max_excitation_error, minima_number)**2
 
 
-def atanc(excitation_error, max_excitation_error):
+def atanc(excitation_error, max_excitation_error, minima_number=5):
     """
     Intensity with arctan(s)/s profile that closely follows sin(s)/s but
     is smooth for s!=0.
@@ -121,11 +121,15 @@ def atanc(excitation_error, max_excitation_error):
     max_excitation_error : float
         The distance at which a reflection becomes extinct
 
+    minima_number : int
+        The minima_number'th minima in the corresponding sinx/x lies at
+        max_excitation_error from 0
+
     Returns
     -------
     intensity : array-like or float
     """
-    fac = np.pi * 5 / np.abs(max_excitation_error)
+    fac = np.pi * minima_number / np.abs(max_excitation_error)
     return np.nan_to_num(
             np.arctan(fac*excitation_error)/(fac*excitation_error),
             nan=1,
@@ -150,8 +154,8 @@ def lorentzian(excitation_error, max_excitation_error):
     intensity_factor : array-like or float
         Vector representing the rel-rod factor for each reflection
 
-    Notes
-    -----
+    References
+    ----------
     [1] L. Palatinus, P. Brázda, M. Jelínek, J. Hrdá, G. Steciuk, M. Klementová, Specifics of the data processing of precession electron diffraction tomography data and their implementation in the program PETS2.0, Acta Crystallogr. Sect. B Struct. Sci. Cryst. Eng. Mater. 75 (2019) 512–522. doi:10.1107/S2052520619007534.
     """
     # in the paper, sigma = pi*thickness.
@@ -187,8 +191,8 @@ def lorentzian_precession(excitation_error, max_excitation_error,
     intensity_factor : array-like or float
         Vector representing the rel-rod factor for each reflection
 
-    Notes
-    -----
+    References
+    ----------
     [1] L. Palatinus, P. Brázda, M. Jelínek, J. Hrdá, G. Steciuk, M. Klementová, Specifics of the data processing of precession electron diffraction tomography data and their implementation in the program PETS2.0, Acta Crystallogr. Sect. B Struct. Sci. Cryst. Eng. Mater. 75 (2019) 512–522. doi:10.1107/S2052520619007534.
     """
     sigma = np.pi/max_excitation_error

From 10f23909655246f93dca599e685782facfd35841 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Wed, 18 Nov 2020 12:22:27 +0100
Subject: [PATCH 05/10] updated tests added some tests

---
 diffsims/generators/diffraction_generator.py  |  7 ++-
 .../test_diffraction_generator.py             | 61 ++++++++++++++++---
 2 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index 13ed8502..e481526a 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/diffsims/generators/diffraction_generator.py
@@ -357,7 +357,8 @@ 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.
@@ -372,6 +373,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
         -------
@@ -399,7 +402,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):
diff --git a/diffsims/tests/test_generators/test_diffraction_generator.py b/diffsims/tests/test_generators/test_diffraction_generator.py
index c1121e0f..8655d442 100644
--- a/diffsims/tests/test_generators/test_diffraction_generator.py
+++ b/diffsims/tests/test_generators/test_diffraction_generator.py
@@ -23,10 +23,14 @@
 from diffsims.generators.diffraction_generator import (
     DiffractionGenerator,
     AtomicDiffractionGenerator,
+    _z_sphere_precession,
+    _shape_factor_precession,
+    _average_excitation_error_precession,
 )
 import diffpy.structure
-from diffsims.utils.shape_factor_models import linear, binary, sinc
-
+from diffsims.utils.shape_factor_models import (linear, binary, sinc, sin2c,
+                                                atanc, lorentzian, 
+                                                lorentzian_precession)
 
 @pytest.fixture(params=[(300)])
 def diffraction_calculator(request):
@@ -82,10 +86,45 @@ def probe(x, out=None, scale=None):
         return v + 0 * x[2].reshape(1, 1, -1)
 
 
+@pytest.mark.parametrize("parameters, expected",
+                         [([0, 1, 0.001, 0.5], -0.00822681491001731),
+                          ([0, np.array([1, 2, 20]), 0.001, 0.5],
+                            np.array([-0.00822681, -0.01545354, 0.02547058])),
+                          ([180, 1, 0.001, 0.5], 0.00922693)])
+def test_z_sphere_precession(parameters, expected):
+    result = _z_sphere_precession(*parameters)
+    assert np.allclose(result, expected)
+
+
+@pytest.mark.parametrize("model", [linear, atanc, sin2c, lorentzian])
+def test_shape_factor_precession(model):
+    z = np.array([-0.1, 0.1])
+    r = np.array([1, 5])
+    _ = _shape_factor_precession(z, r, 0.001, 0.5, model, 0.1)
+
+
+def test_average_excitation_error_precession():
+    z = np.array([-0.1, 0.1])
+    r = np.array([1, 5])
+    _ = _average_excitation_error_precession(z, r, 0.001, 0.5)
+
+
+@pytest.mark.parametrize("model, expected",
+                         [("linear", linear),
+                          ("lorentzian", lorentzian),
+                          (binary, binary)],)
+def test_diffraction_generator_init(model, expected):
+    generator = DiffractionGenerator(300, shape_factor_model=model)
+    assert generator.shape_factor_model == expected
+
+
 class TestDiffractionCalculator:
     def test_init(self, diffraction_calculator: DiffractionGenerator):
-        assert diffraction_calculator.debye_waller_factors == {}
-        _ = DiffractionGenerator(300, 2)
+        assert diffraction_calculator.scattering_params == "lobato"
+        assert diffraction_calculator.precession_angle == 0
+        assert diffraction_calculator.shape_factor_model == linear
+        assert diffraction_calculator.approximate_precession == True
+        assert diffraction_calculator.minimum_intensity == 1e-20
 
     def test_matching_results(self, diffraction_calculator, local_structure):
         diffraction = diffraction_calculator.calculate_ed_data(
@@ -124,10 +163,9 @@ def test_appropriate_intensities(self, diffraction_calculator, local_structure):
         )
         assert np.all(smaller)
 
-    @pytest.mark.parametrize("model", [None, linear, binary, sinc])
-    def test_shape_factor_strings(self, diffraction_calculator, local_structure, model):
+    def test_shape_factor_strings(self, diffraction_calculator, local_structure):
         _ = diffraction_calculator.calculate_ed_data(
-            local_structure, 2, shape_factor_model=model
+            local_structure, 2
         )
 
     def test_shape_factor_custom(self, diffraction_calculator, local_structure):
@@ -135,10 +173,10 @@ def local_excite(excitation_error, maximum_excitation_error, t):
             return (np.sin(t) * excitation_error) / maximum_excitation_error
 
         t1 = diffraction_calculator.calculate_ed_data(
-            local_structure, 2, shape_factor_model=local_excite, t=0.2
+            local_structure, 2, max_excitation_error=0.02
         )
         t2 = diffraction_calculator.calculate_ed_data(
-            local_structure, 2, shape_factor_model=local_excite, t=0.4
+            local_structure, 2, max_excitation_error=0.4
         )
 
         # softly makes sure the two sims are different
@@ -204,6 +242,11 @@ def test_invalid_scattering_params():
     generator = DiffractionGenerator(300, scattering_params=scattering_param)
 
 
+@pytest.mark.xfail(faises=NotImplementedError)
+def test_invalid_shape_model():
+    generator = DiffractionGenerator(300, shape_factor_model="dracula")
+
+
 @pytest.mark.parametrize("shape", [(10, 20), (20, 10)])
 def test_param_check_atomic(shape):
     detector = [np.linspace(-1, 1, s) for s in shape]

From f4d61c4021cb13816b2aa76d564ca5c1284f1eed Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Wed, 18 Nov 2020 14:04:19 +0100
Subject: [PATCH 06/10] added more test coverage, fixed some bugs

---
 diffsims/generators/diffraction_generator.py  | 21 ++++++-----
 .../test_diffraction_generator.py             | 36 ++++++++++++++++---
 2 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index e481526a..27addaa9 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/diffsims/generators/diffraction_generator.py
@@ -289,16 +289,15 @@ def calculate_ed_data(
         r_spot = np.sqrt(np.sum(np.square(cartesian_coordinates[:, :2]), axis=1))
         z_spot = cartesian_coordinates[:, 2]
 
-        if self.precession_angle > 0:
-            if 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,
-                        )
+        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
@@ -317,7 +316,7 @@ def calculate_ed_data(
                     wavelength,
                     self.precession_angle,
                     self.shape_factor_model,
-                    self.max_excitation_error,
+                    max_excitation_error,
                     **self.shape_factor_kwargs,
                     )
         elif self.precession_angle > 0 and self.approximate_precession:
diff --git a/diffsims/tests/test_generators/test_diffraction_generator.py b/diffsims/tests/test_generators/test_diffraction_generator.py
index 8655d442..fae195ca 100644
--- a/diffsims/tests/test_generators/test_diffraction_generator.py
+++ b/diffsims/tests/test_generators/test_diffraction_generator.py
@@ -28,15 +28,27 @@
     _average_excitation_error_precession,
 )
 import diffpy.structure
-from diffsims.utils.shape_factor_models import (linear, binary, sinc, sin2c,
-                                                atanc, lorentzian, 
-                                                lorentzian_precession)
+from diffsims.utils.shape_factor_models import (linear, binary, sin2c,
+                                                atanc, lorentzian)
+
 
 @pytest.fixture(params=[(300)])
 def diffraction_calculator(request):
     return DiffractionGenerator(request.param)
 
 
+@pytest.fixture(scope="module")
+def diffraction_calculator_precession_full():
+    return DiffractionGenerator(300, precession_angle=0.5,
+                                approximate_precession=False)
+
+
+@pytest.fixture(scope="module")
+def diffraction_calculator_precession_simple():
+    return DiffractionGenerator(300, precession_angle=0.5,
+                                approximate_precession=True)
+
+
 @pytest.fixture(params=[(300, [np.linspace(-1, 1, 10)] * 2)])
 def diffraction_calculator_atomic(request):
     return AtomicDiffractionGenerator(*request.param)
@@ -96,7 +108,7 @@ def test_z_sphere_precession(parameters, expected):
     assert np.allclose(result, expected)
 
 
-@pytest.mark.parametrize("model", [linear, atanc, sin2c, lorentzian])
+@pytest.mark.parametrize("model", [binary, linear, atanc, sin2c, lorentzian])
 def test_shape_factor_precession(model):
     z = np.array([-0.1, 0.1])
     r = np.array([1, 5])
@@ -133,6 +145,22 @@ def test_matching_results(self, diffraction_calculator, local_structure):
         assert len(diffraction.indices) == len(diffraction.coordinates)
         assert len(diffraction.coordinates) == len(diffraction.intensities)
 
+    def test_precession_simple(self, diffraction_calculator_precession_simple,
+            local_structure):
+        diffraction = diffraction_calculator_precession_simple.calculate_ed_data(
+            local_structure, reciprocal_radius=5.0,
+        )
+        assert len(diffraction.indices) == len(diffraction.coordinates)
+        assert len(diffraction.coordinates) == len(diffraction.intensities)
+
+    def test_precession_full(self, diffraction_calculator_precession_full,
+            local_structure):
+        diffraction = diffraction_calculator_precession_full.calculate_ed_data(
+            local_structure, reciprocal_radius=5.0,
+        )
+        assert len(diffraction.indices) == len(diffraction.coordinates)
+        assert len(diffraction.coordinates) == len(diffraction.intensities)
+        
     def test_appropriate_scaling(self, diffraction_calculator: DiffractionGenerator):
         """Tests that doubling the unit cell halves the pattern spacing."""
         silicon = make_structure(5)

From 598a95ed64866b5c1cbb945803f353f4ace35915 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Wed, 18 Nov 2020 14:32:31 +0100
Subject: [PATCH 07/10] added custom shape function test

---
 .../test_diffraction_generator.py             | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/diffsims/tests/test_generators/test_diffraction_generator.py b/diffsims/tests/test_generators/test_diffraction_generator.py
index fae195ca..b8e43448 100644
--- a/diffsims/tests/test_generators/test_diffraction_generator.py
+++ b/diffsims/tests/test_generators/test_diffraction_generator.py
@@ -49,6 +49,17 @@ def diffraction_calculator_precession_simple():
                                 approximate_precession=True)
 
 
+def local_excite(excitation_error, maximum_excitation_error, t):
+    return (np.sin(t) * excitation_error) / maximum_excitation_error
+
+
+@pytest.fixture(scope="module")
+def diffraction_calculator_custom():
+    return DiffractionGenerator(300,
+                                shape_factor_model=local_excite,
+                                t=0.2)
+
+
 @pytest.fixture(params=[(300, [np.linspace(-1, 1, 10)] * 2)])
 def diffraction_calculator_atomic(request):
     return AtomicDiffractionGenerator(*request.param)
@@ -161,6 +172,14 @@ def test_precession_full(self, diffraction_calculator_precession_full,
         assert len(diffraction.indices) == len(diffraction.coordinates)
         assert len(diffraction.coordinates) == len(diffraction.intensities)
         
+    def test_custom_shape_func(self, diffraction_calculator_custom,
+            local_structure):
+        diffraction = diffraction_calculator_custom.calculate_ed_data(
+            local_structure, reciprocal_radius=5.0,
+        )
+        assert len(diffraction.indices) == len(diffraction.coordinates)
+        assert len(diffraction.coordinates) == len(diffraction.intensities)
+
     def test_appropriate_scaling(self, diffraction_calculator: DiffractionGenerator):
         """Tests that doubling the unit cell halves the pattern spacing."""
         silicon = make_structure(5)
@@ -197,8 +216,6 @@ def test_shape_factor_strings(self, diffraction_calculator, local_structure):
         )
 
     def test_shape_factor_custom(self, diffraction_calculator, local_structure):
-        def local_excite(excitation_error, maximum_excitation_error, t):
-            return (np.sin(t) * excitation_error) / maximum_excitation_error
 
         t1 = diffraction_calculator.calculate_ed_data(
             local_structure, 2, max_excitation_error=0.02

From d24eb78eb5ebbc9fc8be8351d26a3cb409f0c270 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Wed, 18 Nov 2020 20:54:47 +0100
Subject: [PATCH 08/10] reformatted with black

---
 diffsims/generators/diffraction_generator.py  | 115 ++++++++++--------
 diffsims/generators/library_generator.py      |  12 +-
 diffsims/generators/sphere_mesh_generators.py |   2 +-
 .../test_diffraction_generator.py             |  65 +++++-----
 4 files changed, 102 insertions(+), 92 deletions(-)

diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index 27addaa9..f9c13b10 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/diffsims/generators/diffraction_generator.py
@@ -33,22 +33,22 @@
 )
 from diffsims.utils.fourier_transform import from_recip
 from diffsims.utils.shape_factor_models import (
-        linear,
-        atanc,
-        lorentzian,
-        sinc,
-        sin2c,
-        lorentzian_precession,
-        )
+    linear,
+    atanc,
+    lorentzian,
+    sinc,
+    sin2c,
+    lorentzian_precession,
+)
 
 
 _shape_factor_model_mapping = {
-        "linear": linear,
-        "atanc": atanc,
-        "sinc": sinc,
-        "sin2c": sin2c,
-        "lorentzian": lorentzian,
-        }
+    "linear": linear,
+    "atanc": atanc,
+    "sinc": sinc,
+    "sin2c": sin2c,
+    "lorentzian": lorentzian,
+}
 
 
 def _z_sphere_precession(phi, r_spot, wavelength, theta):
@@ -75,15 +75,17 @@ def _z_sphere_precession(phi, r_spot, wavelength, theta):
         The height of the ewald sphere at the point r in A^-1
     """
     phi = np.deg2rad(phi)
-    r = 1/wavelength
+    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))
+    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):
+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
@@ -120,29 +122,30 @@ def _shape_factor_precession(z_spot, r_spot, wavelength, precession_angle,
     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)
+            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]
+        shf[i] = (1 / (360)) * quad(integrand, 0, 360)[0]
     return shf
 
 
-def _average_excitation_error_precession(z_spot, r_spot,
-                                         wavelength, precession_angle):
+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
+            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]
+        ext[i] = (1 / (360)) * quad(integrand, 0, 360)[0]
     return ext
 
 
@@ -172,7 +175,7 @@ class DiffractionGenerator(object):
         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.    
+        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
@@ -206,7 +209,9 @@ def __init__(
         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]
+                self.shape_factor_model = _shape_factor_model_mapping[
+                    shape_factor_model
+                ]
             else:
                 raise NotImplementedError(
                     f"{shape_factor_model} is not a recognized shape factor "
@@ -293,11 +298,11 @@ def calculate_ed_data(
             # 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,
-                    )
+                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
@@ -311,25 +316,24 @@ def calculate_ed_data(
         # 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,
-                    )
+                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,
-                    )
+                excitation_error,
+                max_excitation_error,
+                r_spot,
+                self.precession_angle,
+            )
         else:
             shape_factor = self.shape_factor_model(
-                excitation_error, max_excitation_error,
-                **self.shape_factor_kwargs
+                excitation_error, max_excitation_error, **self.shape_factor_kwargs
             )
 
         # Calculate diffracted intensities based on a kinematical model.
@@ -356,8 +360,11 @@ def calculate_ed_data(
         )
 
     def calculate_profile_data(
-        self, structure, reciprocal_radius=1.0, minimum_intensity=1e-3,
-        debye_waller_factors={}
+        self,
+        structure,
+        reciprocal_radius=1.0,
+        minimum_intensity=1e-3,
+        debye_waller_factors={},
     ):
         """Calculates a one dimensional diffraction profile for a
         structure.
@@ -469,7 +476,7 @@ def calculate_ed_data(
         dtype="float64",
         ZERO=1e-14,
         mode="kinematic",
-        **kwargs
+        **kwargs,
     ):
         """
         Calculates single electron diffraction image for particular atomic
diff --git a/diffsims/generators/library_generator.py b/diffsims/generators/library_generator.py
index e45b29ee..03c5f6e1 100644
--- a/diffsims/generators/library_generator.py
+++ b/diffsims/generators/library_generator.py
@@ -111,7 +111,7 @@ def get_diffraction_library(
                     rotation=orientation,
                     with_direct_beam=with_direct_beam,
                     max_excitation_error=max_excitation_error,
-                    debye_waller_factors=debye_waller_factors
+                    debye_waller_factors=debye_waller_factors,
                 )
 
                 # Calibrate simulation
@@ -126,11 +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
diff --git a/diffsims/generators/sphere_mesh_generators.py b/diffsims/generators/sphere_mesh_generators.py
index 1e2b289a..5b7e18d7 100644
--- a/diffsims/generators/sphere_mesh_generators.py
+++ b/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, np.pi/2 - phi2]).T)
+    grid = np.rad2deg(np.vstack([phi1, Phi, np.pi / 2 - phi2]).T)
     return grid
diff --git a/diffsims/tests/test_generators/test_diffraction_generator.py b/diffsims/tests/test_generators/test_diffraction_generator.py
index b8e43448..222ab1ba 100644
--- a/diffsims/tests/test_generators/test_diffraction_generator.py
+++ b/diffsims/tests/test_generators/test_diffraction_generator.py
@@ -28,8 +28,7 @@
     _average_excitation_error_precession,
 )
 import diffpy.structure
-from diffsims.utils.shape_factor_models import (linear, binary, sin2c,
-                                                atanc, lorentzian)
+from diffsims.utils.shape_factor_models import linear, binary, sin2c, atanc, lorentzian
 
 
 @pytest.fixture(params=[(300)])
@@ -39,14 +38,12 @@ def diffraction_calculator(request):
 
 @pytest.fixture(scope="module")
 def diffraction_calculator_precession_full():
-    return DiffractionGenerator(300, precession_angle=0.5,
-                                approximate_precession=False)
+    return DiffractionGenerator(300, precession_angle=0.5, approximate_precession=False)
 
 
 @pytest.fixture(scope="module")
 def diffraction_calculator_precession_simple():
-    return DiffractionGenerator(300, precession_angle=0.5,
-                                approximate_precession=True)
+    return DiffractionGenerator(300, precession_angle=0.5, approximate_precession=True)
 
 
 def local_excite(excitation_error, maximum_excitation_error, t):
@@ -55,9 +52,7 @@ def local_excite(excitation_error, maximum_excitation_error, t):
 
 @pytest.fixture(scope="module")
 def diffraction_calculator_custom():
-    return DiffractionGenerator(300,
-                                shape_factor_model=local_excite,
-                                t=0.2)
+    return DiffractionGenerator(300, shape_factor_model=local_excite, t=0.2)
 
 
 @pytest.fixture(params=[(300, [np.linspace(-1, 1, 10)] * 2)])
@@ -109,11 +104,17 @@ def probe(x, out=None, scale=None):
         return v + 0 * x[2].reshape(1, 1, -1)
 
 
-@pytest.mark.parametrize("parameters, expected",
-                         [([0, 1, 0.001, 0.5], -0.00822681491001731),
-                          ([0, np.array([1, 2, 20]), 0.001, 0.5],
-                            np.array([-0.00822681, -0.01545354, 0.02547058])),
-                          ([180, 1, 0.001, 0.5], 0.00922693)])
+@pytest.mark.parametrize(
+    "parameters, expected",
+    [
+        ([0, 1, 0.001, 0.5], -0.00822681491001731),
+        (
+            [0, np.array([1, 2, 20]), 0.001, 0.5],
+            np.array([-0.00822681, -0.01545354, 0.02547058]),
+        ),
+        ([180, 1, 0.001, 0.5], 0.00922693),
+    ],
+)
 def test_z_sphere_precession(parameters, expected):
     result = _z_sphere_precession(*parameters)
     assert np.allclose(result, expected)
@@ -132,10 +133,10 @@ def test_average_excitation_error_precession():
     _ = _average_excitation_error_precession(z, r, 0.001, 0.5)
 
 
-@pytest.mark.parametrize("model, expected",
-                         [("linear", linear),
-                          ("lorentzian", lorentzian),
-                          (binary, binary)],)
+@pytest.mark.parametrize(
+    "model, expected",
+    [("linear", linear), ("lorentzian", lorentzian), (binary, binary)],
+)
 def test_diffraction_generator_init(model, expected):
     generator = DiffractionGenerator(300, shape_factor_model=model)
     assert generator.shape_factor_model == expected
@@ -156,26 +157,30 @@ def test_matching_results(self, diffraction_calculator, local_structure):
         assert len(diffraction.indices) == len(diffraction.coordinates)
         assert len(diffraction.coordinates) == len(diffraction.intensities)
 
-    def test_precession_simple(self, diffraction_calculator_precession_simple,
-            local_structure):
+    def test_precession_simple(
+        self, diffraction_calculator_precession_simple, local_structure
+    ):
         diffraction = diffraction_calculator_precession_simple.calculate_ed_data(
-            local_structure, reciprocal_radius=5.0,
+            local_structure,
+            reciprocal_radius=5.0,
         )
         assert len(diffraction.indices) == len(diffraction.coordinates)
         assert len(diffraction.coordinates) == len(diffraction.intensities)
 
-    def test_precession_full(self, diffraction_calculator_precession_full,
-            local_structure):
+    def test_precession_full(
+        self, diffraction_calculator_precession_full, local_structure
+    ):
         diffraction = diffraction_calculator_precession_full.calculate_ed_data(
-            local_structure, reciprocal_radius=5.0,
+            local_structure,
+            reciprocal_radius=5.0,
         )
         assert len(diffraction.indices) == len(diffraction.coordinates)
         assert len(diffraction.coordinates) == len(diffraction.intensities)
-        
-    def test_custom_shape_func(self, diffraction_calculator_custom,
-            local_structure):
+
+    def test_custom_shape_func(self, diffraction_calculator_custom, local_structure):
         diffraction = diffraction_calculator_custom.calculate_ed_data(
-            local_structure, reciprocal_radius=5.0,
+            local_structure,
+            reciprocal_radius=5.0,
         )
         assert len(diffraction.indices) == len(diffraction.coordinates)
         assert len(diffraction.coordinates) == len(diffraction.intensities)
@@ -211,9 +216,7 @@ def test_appropriate_intensities(self, diffraction_calculator, local_structure):
         assert np.all(smaller)
 
     def test_shape_factor_strings(self, diffraction_calculator, local_structure):
-        _ = diffraction_calculator.calculate_ed_data(
-            local_structure, 2
-        )
+        _ = diffraction_calculator.calculate_ed_data(local_structure, 2)
 
     def test_shape_factor_custom(self, diffraction_calculator, local_structure):
 

From ea7cb65710a76219779c6e18770f5ae4fde68720 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Mon, 23 Nov 2020 16:55:56 +0100
Subject: [PATCH 09/10] implemented review, added to changelog

---
 CHANGELOG.md                                 |  7 ++-
 diffsims/generators/diffraction_generator.py | 61 +++++++++++---------
 diffsims/utils/shape_factor_models.py        | 36 ++++++------
 3 files changed, 59 insertions(+), 45 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c9994e74..5666f417 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index f9c13b10..3d09cc26 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/diffsims/generators/diffraction_generator.py
@@ -51,40 +51,48 @@
 }
 
 
-def _z_sphere_precession(phi, r_spot, wavelength, theta):
+def _z_sphere_precession(theta, r_spot, wavelength, phi):
     """
     Returns the z-coordinate in reciprocal space of the Ewald sphere at
-    distance r from the origin along the x-axis
+    distance r_spot from the origin
 
     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.
+    theta : float
+        The azimuthal angle in degrees
     r_spot : float
-        The distance to the point on the x-axis where we calculate z in A^-1
+        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^-1
-    theta : 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.
     """
-    phi = np.deg2rad(phi)
-    r = 1 / wavelength
     theta = np.deg2rad(theta)
+    r = 1 / wavelength
+    phi = np.deg2rad(phi)
     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)
+        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, precession_angle, function, max_excitation, **kwargs
+    z_spot, r_spot, wavelength, phi, shape_function, max_excitation, **kwargs
 ):
     """
     The rel-rod shape factors for reflections taking into account
@@ -98,34 +106,35 @@ def _shape_factor_precession(
         An array representing the distance of spots from the z-axis in A^-1
     wavelength : float
         The electron wavelength in A
-    precession_angle : float
+    phi : float
         The precession angle in degrees
-    function : callable
+    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
-        Maximum excitation angle to consider if precession_angle = 0. With
-        precession, it is rather a parameter to describe the "width" of the
-        rel-rods.
+        Parameter to describe the "extent" of the rel-rods.
 
     Other parameters
     ----------------
-    ** kwargs: passed directly to function
+    ** kwargs: passed directly to shape_function
 
     Notes
     -----
-    We calculate excitation_error as z_spot - z_sphere so that it is
+    * 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, precession_angle)
-            return function(z_spot_i - z_sph, max_excitation, **kwargs)
+            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]
@@ -175,7 +184,7 @@ class DiffractionGenerator(object):
         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.
+        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
@@ -190,7 +199,7 @@ class DiffractionGenerator(object):
     * 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
+    model defaults to Lorentzian; shape_factor_model is ignored. Only with
     approximate_precession=False the custom shape_factor_model is used.
     """
 
diff --git a/diffsims/utils/shape_factor_models.py b/diffsims/utils/shape_factor_models.py
index c14f4a00..fb867aa7 100644
--- a/diffsims/utils/shape_factor_models.py
+++ b/diffsims/utils/shape_factor_models.py
@@ -81,9 +81,9 @@ def sinc(excitation_error, max_excitation_error, minima_number=5):
     num = np.sin(fac * excitation_error)
     denom = fac * excitation_error
     return np.nan_to_num(
-            np.abs(np.divide(num, denom, out=np.zeros_like(num), where=denom != 0)),
-            nan=1,
-            )
+        np.abs(np.divide(num, denom, out=np.zeros_like(num), where=denom != 0)),
+        nan=1,
+    )
 
 
 def sin2c(excitation_error, max_excitation_error, minima_number=5):
@@ -105,7 +105,7 @@ def sin2c(excitation_error, max_excitation_error, minima_number=5):
     -------
     intensity : array-like or float
     """
-    return sinc(excitation_error, max_excitation_error, minima_number)**2
+    return sinc(excitation_error, max_excitation_error, minima_number) ** 2
 
 
 def atanc(excitation_error, max_excitation_error, minima_number=5):
@@ -131,15 +131,15 @@ def atanc(excitation_error, max_excitation_error, minima_number=5):
     """
     fac = np.pi * minima_number / np.abs(max_excitation_error)
     return np.nan_to_num(
-            np.arctan(fac*excitation_error)/(fac*excitation_error),
-            nan=1,
-            )
+        np.arctan(fac * excitation_error) / (fac * excitation_error),
+        nan=1,
+    )
 
 
 def lorentzian(excitation_error, max_excitation_error):
     """
     Lorentzian intensity profile that should approximate
-    the two-beam rocking curve. After [1].
+    the two-beam rocking curve. This is equation (6) in reference [1].
 
     Parameters
     ----------
@@ -160,16 +160,18 @@ def lorentzian(excitation_error, max_excitation_error):
     """
     # in the paper, sigma = pi*thickness.
     # We assume thickness = 1/max_exitation_error
-    sigma = np.pi/max_excitation_error
-    fac = sigma/(np.pi*(sigma**2*excitation_error**2+1))
+    sigma = np.pi / max_excitation_error
+    fac = sigma / (np.pi * (sigma ** 2 * excitation_error ** 2 + 1))
     return fac
 
 
-def lorentzian_precession(excitation_error, max_excitation_error,
-                          r_spot, precession_angle):
+def lorentzian_precession(
+    excitation_error, max_excitation_error, r_spot, precession_angle
+):
     """
     Intensity profile factor for a precessed beam assuming a Lorentzian
-    intensity profile for the un-precessed beam. After [1].
+    intensity profile for the un-precessed beam. This is equation (10) in
+    reference [1].
 
     Parameters
     ----------
@@ -195,8 +197,8 @@ def lorentzian_precession(excitation_error, max_excitation_error,
     ----------
     [1] L. Palatinus, P. Brázda, M. Jelínek, J. Hrdá, G. Steciuk, M. Klementová, Specifics of the data processing of precession electron diffraction tomography data and their implementation in the program PETS2.0, Acta Crystallogr. Sect. B Struct. Sci. Cryst. Eng. Mater. 75 (2019) 512–522. doi:10.1107/S2052520619007534.
     """
-    sigma = np.pi/max_excitation_error
-    u = sigma**2*(r_spot**2*precession_angle**2 - excitation_error**2)+1
-    z = np.sqrt(u**2 + 4*sigma**2 + excitation_error**2)
-    fac = (sigma/np.pi)*np.sqrt(2*(u+z)/z**2)
+    sigma = np.pi / max_excitation_error
+    u = sigma ** 2 * (r_spot ** 2 * precession_angle ** 2 - excitation_error ** 2) + 1
+    z = np.sqrt(u ** 2 + 4 * sigma ** 2 + excitation_error ** 2)
+    fac = (sigma / np.pi) * np.sqrt(2 * (u + z) / z ** 2)
     return fac

From 301963aef469351e4a2a181565689606cdca3e7f Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Fri, 27 Nov 2020 16:43:47 +0100
Subject: [PATCH 10/10] changed default shape factor to lorentzian

---
 diffsims/generators/diffraction_generator.py                 | 2 +-
 diffsims/tests/test_generators/test_diffraction_generator.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/diffsims/generators/diffraction_generator.py b/diffsims/generators/diffraction_generator.py
index 3d09cc26..d6361bb2 100644
--- a/diffsims/generators/diffraction_generator.py
+++ b/diffsims/generators/diffraction_generator.py
@@ -208,7 +208,7 @@ def __init__(
         accelerating_voltage,
         scattering_params="lobato",
         precession_angle=0,
-        shape_factor_model="linear",
+        shape_factor_model="lorentzian",
         approximate_precession=True,
         minimum_intensity=1e-20,
         **kwargs,
diff --git a/diffsims/tests/test_generators/test_diffraction_generator.py b/diffsims/tests/test_generators/test_diffraction_generator.py
index 222ab1ba..9db28465 100644
--- a/diffsims/tests/test_generators/test_diffraction_generator.py
+++ b/diffsims/tests/test_generators/test_diffraction_generator.py
@@ -146,7 +146,7 @@ class TestDiffractionCalculator:
     def test_init(self, diffraction_calculator: DiffractionGenerator):
         assert diffraction_calculator.scattering_params == "lobato"
         assert diffraction_calculator.precession_angle == 0
-        assert diffraction_calculator.shape_factor_model == linear
+        assert diffraction_calculator.shape_factor_model == lorentzian
         assert diffraction_calculator.approximate_precession == True
         assert diffraction_calculator.minimum_intensity == 1e-20