Merge 9af5f65 into af97040

diffsims/generators/diffraction_generator.py

@@ -36,9 +36,10 @@
     get_vectorized_list_for_atomic_scattering_factors,
     is_lattice_hexagonal,
     get_intensities_params,
-    get_scattering_params_dict
+    get_scattering_params_dict,
 )
 from diffsims.utils.fourier_transform import from_recip
+from diffsims.utils.shape_factor_models import linear, binary, sinc
 
 
 class DiffractionGenerator(object):
@@ -72,15 +73,15 @@ def __init__(
         accelerating_voltage,
         max_excitation_error=None,
         debye_waller_factors={},
-        scattering_params="lobato"
+        scattering_params="lobato",
     ):
         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
-        if scattering_params in ["lobato","xtables"]:
+        if scattering_params in ["lobato", "xtables"]:
             self.scattering_params = scattering_params
         else:
             raise NotImplementedError(
@@ -94,7 +95,7 @@ def calculate_ed_data(
         structure,
         reciprocal_radius,
         rotation=(0, 0, 0),
-        shape_factor_model="linear",
+        shape_factor_model=None,
         max_excitation_error=1e-2,
         with_direct_beam=True,
         **kwargs
@@ -113,9 +114,9 @@ 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 str
+        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. The code provides "linear" and "binary" options accessed with by parsing the associated strings
+            scaling factor. If None defaults to shape_factor_models.linear
         max_excitation_error : float
             the exctinction distance for reflections, in reciprocal Angstroms
         with_direct_beam : bool
@@ -162,14 +163,12 @@ def calculate_ed_data(
         excitation_error = excitation_error[intersection]
         g_hkls = spot_distances[intersection]
 
-        if shape_factor_model == "linear":
-            shape_factor = 1 - (excitation_error / max_excitation_error)
-        elif shape_factor_model == "binary":
-            shape_factor = 1
-        else:
+        if shape_factor_model is not None:
             shape_factor = shape_factor_model(
                 excitation_error, max_excitation_error, **kwargs
             )
+        else:
+            shape_factor = linear(excitation_error, max_excitation_error)
 
         # Calculate diffracted intensities based on a kinematical model.
         intensities = get_kinematical_intensities(

diffsims/generators/rotation_list_generators.py

@@ -68,7 +68,9 @@ def get_list_from_orix(grid, rounding=2):
     rotation_list = z.data.tolist()
     i = 0
     while i < len(rotation_list):
-        rotation_list[i] = tuple(np.round(np.rad2deg(rotation_list[i]), decimals=rounding))
+        rotation_list[i] = tuple(
+            np.round(np.rad2deg(rotation_list[i]), decimals=rounding)
+        )
         i += 1
 
     return rotation_list
@@ -115,9 +117,9 @@ def get_local_grid(resolution=2, center=None, grid_width=10):
     -------
     rotation_list : list of tuples
     """
-    if isinstance(center,tuple):
+    if isinstance(center, tuple):
         z = np.deg2rad(np.asarray(center))
-        center = Rotation.from_euler(z,convention="bunge",direction="crystal2lab")
+        center = Rotation.from_euler(z, convention="bunge", direction="crystal2lab")
 
     orix_grid = get_sample_local(
         resolution=resolution, center=center, grid_width=grid_width
@@ -152,14 +154,16 @@ def get_grid_around_beam_direction(beam_rotation, resolution, angular_range=(0,
     >>> grid = get_grid_around_beam_direction(beam_rotation,1)
     """
     z = np.deg2rad(np.asarray(beam_rotation))
-    beam_rotation = Rotation.from_euler(z,convention="bunge",direction="crystal2lab")
+    beam_rotation = Rotation.from_euler(z, convention="bunge", direction="crystal2lab")
 
-    angles = np.deg2rad(np.arange(start=angular_range[0],stop=angular_range[1],step=resolution))
-    axes   = np.repeat([[0,0,1]],angles.shape[0],axis=0)
-    in_plane_rotation = Rotation.from_neo_euler(AxAngle.from_axes_angles(axes,angles))
+    angles = np.deg2rad(
+        np.arange(start=angular_range[0], stop=angular_range[1], step=resolution)
+    )
+    axes = np.repeat([[0, 0, 1]], angles.shape[0], axis=0)
+    in_plane_rotation = Rotation.from_neo_euler(AxAngle.from_axes_angles(axes, angles))
 
     orix_grid = beam_rotation * in_plane_rotation
-    rotation_list = get_list_from_orix(orix_grid,rounding=2)
+    rotation_list = get_list_from_orix(orix_grid, rounding=2)
     return rotation_list
 
 

diffsims/tests/test_generators/test_diffraction_generator.py

@@ -26,6 +26,7 @@
     AtomicDiffractionGenerator,
 )
 import diffpy.structure
+from diffsims.utils.shape_factor_models import linear, binary, sinc
 
 
 @pytest.fixture(params=[(300)])
@@ -124,19 +125,25 @@ def test_appropriate_intensities(self, diffraction_calculator, local_structure):
         )
         assert np.all(smaller)
 
-    @pytest.mark.parametrize("string", ["linear", "binary"])
-    def test_shape_factor_strings(
-        self, diffraction_calculator, local_structure, string
-    ):
+    @pytest.mark.parametrize("model", [None, linear, binary, sinc])
+    def test_shape_factor_strings(self, diffraction_calculator, local_structure, model):
         _ = diffraction_calculator.calculate_ed_data(
-            local_structure, 2, shape_factor_model=string
+            local_structure, 2, shape_factor_model=model
         )
 
     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
 
-        _ = diffraction_calculator.calculate_ed_data(local_structure, 2,shape_factor_model=local_excite, t=0.2)
+        t1 = diffraction_calculator.calculate_ed_data(
+            local_structure, 2, shape_factor_model=local_excite, t=0.2
+        )
+        t2 = diffraction_calculator.calculate_ed_data(
+            local_structure, 2, shape_factor_model=local_excite, t=0.4
+        )
+
+        # softly makes sure the two sims are different
+        assert np.sum(t1.intensities) != np.sum(t2.intensities)
 
     def test_calculate_profile_class(self, local_structure, diffraction_calculator):
         # tests the non-hexagonal (cubic) case

diffsims/tests/test_generators/test_rotation_list_generator.py

@@ -30,7 +30,7 @@
     "grid",
     [
         pytest.param(
-            get_local_grid(resolution=30, center=(0,1,0), grid_width=35),
+            get_local_grid(resolution=30, center=(0, 1, 0), grid_width=35),
             marks=pytest.mark.xfail(reason="Downstream bug"),
         ),
         get_fundamental_zone_grid(space_group=20, resolution=20),
@@ -41,12 +41,16 @@ def test_get_grid(grid):
     assert len(grid) > 0
     assert isinstance(grid[0], tuple)
 
+
 def test_get_grid_around_beam_direction():
-    grid = get_grid_around_beam_direction((0,90,0),resolution=2,angular_range=(0,9))
+    grid = get_grid_around_beam_direction(
+        (0, 90, 0), resolution=2, angular_range=(0, 9)
+    )
     assert isinstance(grid, list)
     assert isinstance(grid[0], tuple)
-    assert len(grid) == 5 # should have 0,2,4,6 and 8
-    assert np.allclose([x[1] for x in grid],90) #taking z to y
+    assert len(grid) == 5  # should have 0,2,4,6 and 8
+    assert np.allclose([x[1] for x in grid], 90)  # taking z to y
+
 
 @pytest.mark.parametrize(
     "crystal_system",

diffsims/utils/shape_factor_models.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+# Copyright 2017-2020 The diffsims developers
+#
+# This file is part of diffsims.
+#
+# diffsims is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# diffsims is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with diffsims.  If not, see <http://www.gnu.org/licenses/>.
+
+import numpy as np
+
+
+def binary(excitation_error, max_excitation_error):
+    """
+    Returns a unit intensity for all reflections
+
+    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
+    -------
+    intensities : array-like or float
+    """
+    return 1
+
+
+def linear(excitation_error, max_excitation_error):
+    """
+    Returns an intensity linearly scaled with by the excitation error
+
+    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
+    -------
+    intensities : array-like or float
+    """
+
+    return 1 - excitation_error / max_excitation_error
+
+
+def sinc(excitation_error, max_excitation_error, minima_number=5):
+    """
+    Returns an intensity with a sinc profile
+
+    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
+    """
+
+    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))