Bugfix: Fix intensity scaling with zero beam

diffsims/generators/simulation_generator.py

@@ -57,6 +57,13 @@ class SimulationGenerator:
     A class for generating kinematic diffraction simulations.
     """
 
+    def __repr__(self):
+        return (
+            f"SimulationGenerator(accelerating_voltage={self.accelerating_voltage}, "
+            f"scattering_params={self.scattering_params}, "
+            f"approximate_precession={self.approximate_precession})"
+        )
+
     def __init__(
         self,
         accelerating_voltage: float = 200,
@@ -67,7 +74,7 @@ def __init__(
         minimum_intensity: float = 1e-20,
         **kwargs,
     ):
-        self.wavelength = get_electron_wavelength(accelerating_voltage)
+        self.accelerating_voltage = accelerating_voltage
         self.precession_angle = np.abs(precession_angle)
         self.approximate_precession = approximate_precession
         if isinstance(shape_factor_model, str):
@@ -94,6 +101,10 @@ def __init__(
                 "implementations.".format(scattering_params)
             )
 
+    @property
+    def wavelength(self):
+        return get_electron_wavelength(self.accelerating_voltage)
+
     def calculate_ed_data(
         self,
         phase: Union[Phase, Sequence[Phase]],
@@ -126,7 +137,7 @@ def calculate_ed_data(
             The cut-off for geometric excitation error in the z-direction
             in units of reciprocal Angstroms. Spots with a larger distance
             from the Ewald sphere are removed from the pattern.
-            Related to the extinction distance and roungly equal to 1/thickness.
+            Related to the extinction distance and roughly equal to 1/thickness.
         shape_factor_width
             Determines the width of the reciprocal rel-rod, for fine-grained
             control. If not set will be set equal to max_excitation_error.
@@ -158,25 +169,30 @@ def calculate_ed_data(
             phase_vectors = []
             for rot in rotate.to_matrix():
                 # Calculate the reciprocal lattice vectors that intersect the Ewald sphere.
-                intersected_vectors, shape_factor = self.get_intersecting_reflections(
+                (
+                    intersected_vectors,
+                    hkl,
+                    shape_factor,
+                ) = self.get_intersecting_reflections(
                     recip,
                     rot,
                     wavelength,
                     max_excitation_error,
                     shape_factor_width=shape_factor_width,
+                    with_direct_beam=with_direct_beam,
                 )
 
                 # Calculate diffracted intensities based on a kinematic model.
                 intensities = get_kinematical_intensities(
                     p.structure,
-                    intersected_vectors.hkl,
+                    hkl,
                     intersected_vectors.gspacing,
                     prefactor=shape_factor,
                     scattering_params=self.scattering_params,
                     debye_waller_factors=debye_waller_factors,
                 )
 
-                # Threshold peaks included in simulation as factor of maximum intensity.
+                # Threshold peaks included in simulation as factor of zero beam intensity.
                 peak_mask = intensities > np.max(intensities) * self.minimum_intensity
                 intensities = intensities[peak_mask]
                 intersected_vectors = intersected_vectors[peak_mask]
@@ -232,7 +248,7 @@ def calculate_diffraction1d(
             recip_latt, reciprocal_radius
         )
 
-        ##spot_indicies is a numpy.array of the hkls allowd in the recip radius
+        ##spot_indicies is a numpy.array of the hkls allowed in the recip radius
         g_indices, multiplicities, g_hkls = get_intensities_params(
             recip_latt, reciprocal_radius
         )
@@ -284,6 +300,7 @@ def calculate_diffraction1d(
             intensities=intensities,
             hkl=hkls,
             reciprocal_radius=reciprocal_radius,
+            wavelength=self.wavelength,
         )
 
     def get_intersecting_reflections(
@@ -293,6 +310,7 @@ def get_intersecting_reflections(
         wavelength: float,
         max_excitation_error: float,
         shape_factor_width: float = None,
+        with_direct_beam: bool = True,
     ):
         """Calculates the reciprocal lattice vectors that intersect the Ewald sphere.
 
@@ -313,11 +331,18 @@ def get_intersecting_reflections(
             Determines the width of the reciprocal rel-rod, for fine-grained
             control. If not set will be set equal to max_excitation_error.
         """
+        initial_hkl = recip.hkl
         rotated_vectors = recip.rotate_from_matrix(rot)
+
+        if with_direct_beam:
+            rotated_vectors = np.vstack([rotated_vectors.data, [0, 0, 0]])
+            initial_hkl = np.vstack([initial_hkl, [0, 0, 0]])
+        else:
+            rotated_vectors = rotated_vectors.data
         # Identify the excitation errors of all points (distance from point to Ewald sphere)
         r_sphere = 1 / wavelength
-        r_spot = np.sqrt(np.sum(np.square(rotated_vectors.data[:, :2]), axis=1))
-        z_spot = rotated_vectors.data[:, 2]
+        r_spot = np.sqrt(np.sum(np.square(rotated_vectors[:, :2]), axis=1))
+        z_spot = rotated_vectors[:, 2]
 
         z_sphere = -np.sqrt(r_sphere**2 - r_spot**2) + r_sphere
         excitation_error = z_sphere - z_spot
@@ -339,8 +364,12 @@ def get_intersecting_reflections(
 
         # select these reflections
         intersected_vectors = rotated_vectors[intersection]
+        intersected_vectors = ReciprocalLatticeVector(
+            phase=recip.phase, xyz=intersected_vectors
+        )
         excitation_error = excitation_error[intersection]
         r_spot = r_spot[intersection]
+        hkl = initial_hkl[intersection]
 
         if shape_factor_width is None:
             shape_factor_width = max_excitation_error
@@ -367,4 +396,4 @@ def get_intersecting_reflections(
                     shape_factor_width,
                     **self.shape_factor_kwargs,
                 )
-        return intersected_vectors, shape_factor
+        return intersected_vectors, hkl, shape_factor

diffsims/simulations/simulation1d.py

@@ -16,17 +16,12 @@
 # You should have received a copy of the GNU General Public License
 # along with diffsims.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Union, Sequence, TYPE_CHECKING, Any
-import copy
+from typing import TYPE_CHECKING
 
 import numpy as np
 import matplotlib.pyplot as plt
 from orix.crystal_map import Phase
-from orix.quaternion import Rotation
-from orix.vector import Vector3d
-
-from diffsims.crystallography.reciprocal_lattice_vector import ReciprocalLatticeVector
-from diffsims.pattern.detector_functions import add_shot_and_point_spread
+from diffsims.utils.sim_utils import get_electron_wavelength
 
 # to avoid circular imports
 if TYPE_CHECKING:  # pragma: no cover
@@ -43,6 +38,7 @@ def __init__(
         intensities: np.ndarray,
         hkl: np.ndarray,
         reciprocal_radius: float,
+        wavelength: float,
     ):
         """Initializes the DiffractionSimulation object with data values for
         the coordinates, indices, intensities, calibration and offset.
@@ -52,19 +48,29 @@ def __init__(
         phase
             The phase of the simulation
         reciprocal_spacing
-            The spacing of the reciprocal lattice vectors
+            The spacing of the reciprocal lattice vectors in A^-1
         intensities
             The intensities of the diffraction spots
         hkl
             The hkl indices of the diffraction spots
         reciprocal_radius
             The radius which the reciprocal lattice spacings are plotted out to
+        wavelength
+            The wavelength of the beam in A^-1
         """
         self.phase = phase
         self.reciprocal_spacing = reciprocal_spacing
         self.intensities = intensities
         self.hkl = hkl
         self.reciprocal_radius = reciprocal_radius
+        self.wavelength = wavelength
+
+    def __repr__(self):
+        return f"Simulation1D(name: {self.phase.name}, wavelength: {self.wavelength})"
+
+    @property
+    def theta(self):
+        return np.arctan2(np.array(self.reciprocal_spacing), 1 / self.wavelength)
 
     def plot(self, ax=None, annotate_peaks=False, fontsize=12, with_labels=True):
         """Plots the 1D diffraction pattern,"""

diffsims/simulations/simulation2d.py

@@ -194,7 +194,7 @@ def __next__(self):
                 coords = self.coordinates[self.phase_index]
             else:
                 coords = self.coordinates
-            if self.has_multiple_vectors:
+            if self.has_multiple_rotations:
                 coords = coords[self.rotation_index]
             else:
                 coords = coords
@@ -374,17 +374,17 @@ def has_multiple_phases(self):
         return self.num_phases > 1
 
     @property
-    def has_multiple_vectors(self):
-        """Returns True if the simulation has multiple vectors"""
-        return self.coordinates.size > 1
+    def has_multiple_rotations(self):
+        """Returns True if the simulation has multiple rotations"""
+        return self.rotations.size > 1
 
     def get_current_coordinates(self):
         """Returns the coordinates of the current phase and rotation"""
         if self.has_multiple_phases:
             return copy.deepcopy(
                 self.coordinates[self.phase_index][self.rotation_index]
             )
-        elif not self.has_multiple_phases and self.has_multiple_vectors:
+        elif not self.has_multiple_phases and self.has_multiple_rotations:
             return copy.deepcopy(self.coordinates[self.rotation_index])
         else:
             return copy.deepcopy(self.coordinates)

diffsims/tests/generators/test_simulation_generator.py

@@ -127,7 +127,7 @@ def test_matching_results(
         diffraction = diffraction_calculator.calculate_ed_data(
             local_structure, reciprocal_radius=5.0
         )
-        assert diffraction.coordinates.size == 72
+        assert diffraction.coordinates.size == 69
 
     def test_precession_simple(
         self, diffraction_calculator_precession_simple, local_structure
@@ -136,7 +136,7 @@ def test_precession_simple(
             local_structure,
             reciprocal_radius=5.0,
         )
-        assert diffraction.coordinates.size == 252
+        assert diffraction.coordinates.size == 249
 
     def test_precession_full(
         self, diffraction_calculator_precession_full, local_structure
@@ -145,7 +145,7 @@ def test_precession_full(
             local_structure,
             reciprocal_radius=5.0,
         )
-        assert diffraction.coordinates.size == 252
+        assert diffraction.coordinates.size == 249
 
     def test_custom_shape_func(self, diffraction_calculator_custom, local_structure):
         diffraction = diffraction_calculator_custom.calculate_ed_data(

diffsims/tests/simulations/test_simulation1d.py

@@ -40,6 +40,7 @@ def simulation1d(self):
             reciprocal_spacing=magnitudes,
             intensities=inten,
             reciprocal_radius=recip,
+            wavelength=0.025,
         )
 
     def test_init(self, simulation1d):