From 302ecfa9278b94448ee7de9afe1b818d09e5e6a2 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Tue, 10 Dec 2024 16:20:31 +0100
Subject: [PATCH 01/17] Subclass Phase for faster initialization

---
 .../crystallography/_diffracting_vector.py    | 35 ++++++++++++++++---
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/diffsims/crystallography/_diffracting_vector.py b/diffsims/crystallography/_diffracting_vector.py
index be9a61f3..46e6c7ea 100644
--- a/diffsims/crystallography/_diffracting_vector.py
+++ b/diffsims/crystallography/_diffracting_vector.py
@@ -18,10 +18,38 @@
 
 from diffsims.crystallography import ReciprocalLatticeVector
 import numpy as np
+from diffpy.structure import Structure
+from orix.crystal_map import Phase
 from orix.vector.miller import _transform_space
 from orix.quaternion import Rotation
 
 
+class RotatedPhase(Phase):
+    """Helper class to speed up rotating the basis of a Phase object.
+    The speedup comes from avoiding a deepcopy of the phase.
+
+    Parameters
+    ----------
+    phase : orix.crystal_map.Phase
+        A phase with a crystal lattice and symmetry.
+    rotation : orix.quaternion.Rotation
+        Rotation to apply to the basis of the phase
+    """
+
+    def __init__(self, phase: Phase, rotation: Rotation):
+        self._structure = Structure(phase.structure)
+        self.name = phase.name
+        self.space_group = phase.space_group
+        self.point_group = phase.point_group
+        self.color = phase.color
+
+        # Perform basis rotation
+        br = self.structure.lattice.baserot
+        # In case the base rotation is set already
+        new_br = br @ rotation.to_matrix().squeeze()
+        self.structure.lattice.setLatPar(baserot=new_br)
+
+
 class DiffractingVector(ReciprocalLatticeVector):
     r"""Reciprocal lattice vectors :math:`(hkl)` for use in electron
     diffraction analysis and simulation.
@@ -151,11 +179,8 @@ def rotate_with_basis(self, rotation):
         if rotation.size != 1:
             raise ValueError("Rotation must be a single rotation")
         # rotate basis
-        new_phase = self.phase.deepcopy()
-        br = new_phase.structure.lattice.baserot
-        # In case the base rotation is set already
-        new_br = br @ rotation.to_matrix().squeeze()
-        new_phase.structure.lattice.setLatPar(baserot=new_br)
+        new_phase = RotatedPhase(self.phase, rotation)
+
         # rotate vectors
         vecs = ~rotation * self.to_miller()
         return ReciprocalLatticeVector(new_phase, xyz=vecs.data)

From a06f2e6863b2a3c7dd28d237594ce9a642a3a518 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Thu, 12 Dec 2024 11:32:14 +0100
Subject: [PATCH 02/17] Rotate Ewald's sphere instead of lattice

---
 diffsims/generators/simulation_generator.py | 44 +++++++++++++--------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index 0a701e3c..06cf7a6f 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -335,7 +335,7 @@ def calculate_diffraction1d(
     def get_intersecting_reflections(
         self,
         recip: DiffractingVector,
-        rot: np.ndarray,
+        rot: Rotation,
         wavelength: float,
         max_excitation_error: float,
         shape_factor_width: float = None,
@@ -361,19 +361,30 @@ def get_intersecting_reflections(
             control. If not set will be set equal to max_excitation_error.
         """
         initial_hkl = recip.hkl
-        rotated_vectors = recip.rotate_with_basis(rotation=rot)
-        rotated_phase = rotated_vectors.phase
-        rotated_vectors = rotated_vectors.data
-        if with_direct_beam:
-            rotated_vectors = np.vstack([rotated_vectors.data, [0, 0, 0]])
-            initial_hkl = np.vstack([initial_hkl, [0, 0, 0]])
+        # rotated_vectors = recip.rotate_with_basis(rotation=rot)
+        # rotated_phase = rotated_vectors.phase
+        # rotated_vectors = rotated_vectors.data
+        # if with_direct_beam:
+        #     rotated_vectors = np.vstack([rotated_vectors.data, [0, 0, 0]])
+        #     initial_hkl = np.vstack([initial_hkl, [0, 0, 0]])
         # 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[:, :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
+        # Instead of rotating vectors, rotate Ewald's sphere to find intersections.
+        # Only rotate the intersecting vectors.
+        # Using notation from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
+        r = 1 / wavelength
+        u = rot.to_matrix().squeeze() @ np.array([0, 0, 1])
+        c = r * u
+        o = recip.data
+
+        diff = o - c
+        dot = np.dot(u, diff.T)
+        nabla = dot**2 - np.linalg.norm(diff, axis=1)**2 + r**2
+        # We know the reflections are all going to intersect twice; no need to look at the sign of nabla
+        sqrt_nabla = np.sqrt(nabla)
+        d = np.min([-dot - sqrt_nabla, -dot + sqrt_nabla], axis=0) 
+
+        excitation_error = d
 
         # determine the pre-selection reflections
         if self.precession_angle == 0:
@@ -391,13 +402,14 @@ def get_intersecting_reflections(
             )
 
         # select these reflections
-        intersected_vectors = rotated_vectors[intersection]
+        intersected_vectors = ~rot * (recip[intersection].to_miller())
+        from diffsims.crystallography._diffracting_vector import RotatedPhase
         intersected_vectors = DiffractingVector(
-            phase=rotated_phase,
-            xyz=intersected_vectors,
+            phase=RotatedPhase(recip.phase, rot),
+            xyz=intersected_vectors.data,
         )
         excitation_error = excitation_error[intersection]
-        r_spot = r_spot[intersection]
+        # r_spot = r_spot[intersection]
         hkl = initial_hkl[intersection]
 
         if shape_factor_width is None:

From ebe79c610502968ecd37a2386e85d88b13ec9ee2 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Thu, 12 Dec 2024 15:51:26 +0100
Subject: [PATCH 03/17] Support precession

---
 diffsims/generators/simulation_generator.py | 78 +++++++++++++++------
 1 file changed, 56 insertions(+), 22 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index 06cf7a6f..68fba270 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -196,7 +196,9 @@ def calculate_diffraction2d(
                 include_zero_vector=with_direct_beam,
             )
             phase_vectors = []
-            for rot in rotate:
+            recip_hkl = recip.hkl # Avoid re-calculating in get_intersecting_reflections
+            from tqdm import tqdm
+            for rot in tqdm(rotate, total=rotate.size):
                 # Calculate the reciprocal lattice vectors that intersect the Ewald sphere.
                 (
                     intersected_vectors,
@@ -204,6 +206,7 @@ def calculate_diffraction2d(
                     shape_factor,
                 ) = self.get_intersecting_reflections(
                     recip,
+                    recip_hkl,
                     rot,
                     wavelength,
                     max_excitation_error,
@@ -335,6 +338,7 @@ def calculate_diffraction1d(
     def get_intersecting_reflections(
         self,
         recip: DiffractingVector,
+        recip_hkl: np.ndarray,
         rot: Rotation,
         wavelength: float,
         max_excitation_error: float,
@@ -360,13 +364,9 @@ 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_with_basis(rotation=rot)
-        # rotated_phase = rotated_vectors.phase
-        # rotated_vectors = rotated_vectors.data
-        # if with_direct_beam:
-        #     rotated_vectors = np.vstack([rotated_vectors.data, [0, 0, 0]])
-        #     initial_hkl = np.vstack([initial_hkl, [0, 0, 0]])
+        if with_direct_beam:
+            recip_hkl = np.vstack([recip_hkl, [0, 0, 0]])
+
         # Identify the excitation errors of all points (distance from point to Ewald sphere)
 
         # Instead of rotating vectors, rotate Ewald's sphere to find intersections.
@@ -376,13 +376,17 @@ def get_intersecting_reflections(
         u = rot.to_matrix().squeeze() @ np.array([0, 0, 1])
         c = r * u
         o = recip.data
+        if with_direct_beam:
+            o = np.vstack([o, [0, 0, 0]])
 
         diff = o - c
         dot = np.dot(u, diff.T)
         nabla = dot**2 - np.linalg.norm(diff, axis=1)**2 + r**2
-        # We know the reflections are all going to intersect twice; no need to look at the sign of nabla
+        # We know the reflections are all going to intersect twice
+        # Therefore, no need to look at the sign of nabla
+        # We also know only the smaller root is important, i.e. -sqrt(nabla)
         sqrt_nabla = np.sqrt(nabla)
-        d = np.min([-dot - sqrt_nabla, -dot + sqrt_nabla], axis=0) 
+        d = -dot - sqrt_nabla
 
         excitation_error = d
 
@@ -390,16 +394,46 @@ def get_intersecting_reflections(
         if self.precession_angle == 0:
             intersection = np.abs(excitation_error) < max_excitation_error
         else:
-            # only consider points that intersect the ewald sphere at some point
-            # the center point of the sphere
-            P_z = r_sphere * np.cos(np.deg2rad(self.precession_angle))
-            P_t = r_sphere * np.sin(np.deg2rad(self.precession_angle))
-            # the extremes of the ewald sphere
-            z_surf_up = P_z - np.sqrt(r_sphere**2 - (r_spot + P_t) ** 2)
-            z_surf_do = P_z - np.sqrt(r_sphere**2 - (r_spot - P_t) ** 2)
-            intersection = (z_spot - max_excitation_error <= z_surf_up) & (
-                z_spot + max_excitation_error >= z_surf_do
-            )
+            # Using the following script to find equations for upper and lower bounds for precessing Ewald's sphere
+            """
+            import sympy
+            import numpy as np
+
+            a = sympy.Symbol("a") # Precession angle
+            r = sympy.Symbol("r") # Ewald's sphere radius
+            rho, z = sympy.symbols("rho z") # cylindrical coordinates of reflection
+
+            rot = lambda ang: np.asarray([[sympy.cos(ang), -sympy.sin(ang)],[sympy.sin(ang), sympy.cos(ang)]])
+
+            u = np.asarray([0, 1])
+            c = r * u
+            cl = rot(a) @ c
+            cr = rot(-a) @ c
+            o = np.asarray([rho, z])
+
+            def get_d(_c):
+                diff = o - _c
+                dot = np.dot(u, diff)
+                nabla = dot**2 - sum(i**2 for i in diff) + r**2
+                sqrt_nabla = nabla**0.5
+                return  -dot - sqrt_nabla
+
+            d = get_d(c)
+            d_upper = get_d(cl)
+            d_lower = get_d(cr)
+
+            print(d.simplify())             # r - z - (r**2 - rho**2)**0.5
+            print((d_upper - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) + rho)**2)**0.5
+            print((d_lower - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) - rho)**2)**0.5
+            """
+            # In the above script, d is the same as before.
+            # We need the distance of the reflections from the incident beam, i.e. the cylindrical coordinate rho
+            rho = np.linalg.norm(np.dot(o, u) * u - o) # using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation
+            a = np.deg2rad(self.precession_angle)
+            upper = r*np.cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*np.sin(a) + rho)**2)**0.5
+            lower = r*np.cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*np.sin(a) - rho)**2)**0.5
+            intersection = (d < upper + excitation_error) & (d > lower - excitation_error)
+
 
         # select these reflections
         intersected_vectors = ~rot * (recip[intersection].to_miller())
@@ -409,8 +443,8 @@ def get_intersecting_reflections(
             xyz=intersected_vectors.data,
         )
         excitation_error = excitation_error[intersection]
-        # r_spot = r_spot[intersection]
-        hkl = initial_hkl[intersection]
+        hkl = recip_hkl[intersection]
+        r_spot = intersected_vectors.norm
 
         if shape_factor_width is None:
             shape_factor_width = max_excitation_error

From ce28dc09b169c7ee6d17587190fc92200d5932c5 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Thu, 12 Dec 2024 17:09:30 +0100
Subject: [PATCH 04/17] Working precession, need to test for correctness

---
 diffsims/generators/simulation_generator.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index 68fba270..83b53fb1 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -381,7 +381,7 @@ def get_intersecting_reflections(
 
         diff = o - c
         dot = np.dot(u, diff.T)
-        nabla = dot**2 - np.linalg.norm(diff, axis=1)**2 + r**2
+        nabla = dot**2 - np.sum(diff**2, axis=1) + r**2
         # We know the reflections are all going to intersect twice
         # Therefore, no need to look at the sign of nabla
         # We also know only the smaller root is important, i.e. -sqrt(nabla)
@@ -428,12 +428,12 @@ def get_d(_c):
             """
             # In the above script, d is the same as before.
             # We need the distance of the reflections from the incident beam, i.e. the cylindrical coordinate rho
-            rho = np.linalg.norm(np.dot(o, u) * u - o) # using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation
+            rho = np.linalg.norm(np.dot(o, u)[:, np.newaxis] * u - o, axis=1) # using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation
             a = np.deg2rad(self.precession_angle)
-            upper = r*np.cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*np.sin(a) + rho)**2)**0.5
-            lower = r*np.cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*np.sin(a) - rho)**2)**0.5
-            intersection = (d < upper + excitation_error) & (d > lower - excitation_error)
-
+            first_half = r*np.cos(a) - r + (r**2 - rho**2)**0.5
+            upper = first_half - (r**2 - (r*np.sin(a) + rho)**2)**0.5
+            lower = first_half - (r**2 - (r*np.sin(a) - rho)**2)**0.5
+            intersection = (d < (upper + max_excitation_error)) & (d > (lower - max_excitation_error))
 
         # select these reflections
         intersected_vectors = ~rot * (recip[intersection].to_miller())

From 1780e3831fb0afa7fad4d6999e0d376aa2ade230 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Fri, 13 Dec 2024 06:52:50 +0100
Subject: [PATCH 05/17] Support direct beam

---
 diffsims/generators/simulation_generator.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index 83b53fb1..c47920df 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -364,9 +364,6 @@ 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.
         """
-        if with_direct_beam:
-            recip_hkl = np.vstack([recip_hkl, [0, 0, 0]])
-
         # Identify the excitation errors of all points (distance from point to Ewald sphere)
 
         # Instead of rotating vectors, rotate Ewald's sphere to find intersections.
@@ -378,6 +375,9 @@ def get_intersecting_reflections(
         o = recip.data
         if with_direct_beam:
             o = np.vstack([o, [0, 0, 0]])
+            recip_hkl = np.vstack([recip_hkl, [0, 0, 0]])
+            # Modify the data directly, as `data` is a property
+            recip._data = np.vstack([recip._data, [0, 0, 0]])
 
         diff = o - c
         dot = np.dot(u, diff.T)

From f931d7335c9176ffd0d979e72b62eaf7e2ae7e65 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Fri, 13 Dec 2024 07:13:06 +0100
Subject: [PATCH 06/17] Fix formatting

---
 diffsims/generators/simulation_generator.py | 22 +++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index c47920df..9523beed 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -24,7 +24,7 @@
 from orix.quaternion import Rotation
 from orix.crystal_map import Phase
 
-from diffsims.crystallography._diffracting_vector import DiffractingVector
+from diffsims.crystallography._diffracting_vector import DiffractingVector, RotatedPhase
 from diffsims.utils.shape_factor_models import (
     linear,
     atanc,
@@ -196,9 +196,9 @@ def calculate_diffraction2d(
                 include_zero_vector=with_direct_beam,
             )
             phase_vectors = []
-            recip_hkl = recip.hkl # Avoid re-calculating in get_intersecting_reflections
-            from tqdm import tqdm
-            for rot in tqdm(rotate, total=rotate.size):
+            # Avoid re-calculating in get_intersecting_reflections
+            recip_hkl = recip.hkl
+            for rot in rotate:
                 # Calculate the reciprocal lattice vectors that intersect the Ewald sphere.
                 (
                     intersected_vectors,
@@ -428,16 +428,18 @@ def get_d(_c):
             """
             # In the above script, d is the same as before.
             # We need the distance of the reflections from the incident beam, i.e. the cylindrical coordinate rho
-            rho = np.linalg.norm(np.dot(o, u)[:, np.newaxis] * u - o, axis=1) # using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation
+            # (using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation):
+            rho = np.linalg.norm(np.dot(o, u)[:, np.newaxis] * u - o, axis=1)
             a = np.deg2rad(self.precession_angle)
-            first_half = r*np.cos(a) - r + (r**2 - rho**2)**0.5
-            upper = first_half - (r**2 - (r*np.sin(a) + rho)**2)**0.5
-            lower = first_half - (r**2 - (r*np.sin(a) - rho)**2)**0.5
-            intersection = (d < (upper + max_excitation_error)) & (d > (lower - max_excitation_error))
+            first_half = r * np.cos(a) - r + (r**2 - rho**2) ** 0.5
+            upper = first_half - (r**2 - (r * np.sin(a) + rho) ** 2) ** 0.5
+            lower = first_half - (r**2 - (r * np.sin(a) - rho) ** 2) ** 0.5
+            intersection = (d < (upper + max_excitation_error)) & (
+                d > (lower - max_excitation_error)
+            )
 
         # select these reflections
         intersected_vectors = ~rot * (recip[intersection].to_miller())
-        from diffsims.crystallography._diffracting_vector import RotatedPhase
         intersected_vectors = DiffractingVector(
             phase=RotatedPhase(recip.phase, rot),
             xyz=intersected_vectors.data,

From bccd856a4147f3fa39b9b3fd936ba361f1562fb7 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Fri, 13 Dec 2024 09:29:54 +0100
Subject: [PATCH 07/17] Fix bug where direct beam is added twice

---
 diffsims/generators/simulation_generator.py   |  6 ++---
 .../generators/test_simulation_generator.py   | 25 ++++++++++++++++---
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index 9523beed..d0420708 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -211,7 +211,7 @@ def calculate_diffraction2d(
                     wavelength,
                     max_excitation_error,
                     shape_factor_width=shape_factor_width,
-                    with_direct_beam=with_direct_beam,
+                    with_direct_beam=False,  # Already handled in DiffractingVector.from_min_dspacing
                 )
 
                 # Calculate diffracted intensities based on a kinematic model.
@@ -223,7 +223,6 @@ def calculate_diffraction2d(
                     scattering_params=self.scattering_params,
                     debye_waller_factors=debye_waller_factors,
                 )
-
                 # Threshold peaks included in simulation as factor of zero beam intensity.
                 peak_mask = intensities > np.max(intensities) * self.minimum_intensity
                 intensities = intensities[peak_mask]
@@ -375,9 +374,9 @@ def get_intersecting_reflections(
         o = recip.data
         if with_direct_beam:
             o = np.vstack([o, [0, 0, 0]])
-            recip_hkl = np.vstack([recip_hkl, [0, 0, 0]])
             # Modify the data directly, as `data` is a property
             recip._data = np.vstack([recip._data, [0, 0, 0]])
+            recip_hkl = np.vstack([recip_hkl, [0, 0, 0]])
 
         diff = o - c
         dot = np.dot(u, diff.T)
@@ -437,7 +436,6 @@ def get_d(_c):
             intersection = (d < (upper + max_excitation_error)) & (
                 d > (lower - max_excitation_error)
             )
-
         # select these reflections
         intersected_vectors = ~rot * (recip[intersection].to_miller())
         intersected_vectors = DiffractingVector(
diff --git a/diffsims/tests/generators/test_simulation_generator.py b/diffsims/tests/generators/test_simulation_generator.py
index 42915baa..67f54723 100644
--- a/diffsims/tests/generators/test_simulation_generator.py
+++ b/diffsims/tests/generators/test_simulation_generator.py
@@ -133,13 +133,30 @@ def test_repr(self, diffraction_calculator):
             "approximate_precession=True)"
         )
 
+    @pytest.mark.parametrize("with_direct_beam", [False, True])
+    def test_single_direct_beam(
+        self,
+        diffraction_calculator: SimulationGenerator,
+        local_structure,
+        with_direct_beam,
+    ):
+        diffraction = diffraction_calculator.calculate_diffraction2d(
+            local_structure, reciprocal_radius=5.0, with_direct_beam=with_direct_beam
+        )
+        # Check that there is only one direct beam
+        num_direct_beam = 0
+        for hkl in diffraction.coordinates.hkl:
+            hkl = tuple(hkl.flatten().tolist())
+            num_direct_beam += hkl == (0, 0, 0)
+        assert num_direct_beam == with_direct_beam
+
     def test_matching_results(
         self, diffraction_calculator: SimulationGenerator, local_structure
     ):
         diffraction = diffraction_calculator.calculate_diffraction2d(
             local_structure, reciprocal_radius=5.0
         )
-        assert diffraction.coordinates.size == 70
+        assert diffraction.coordinates.size == 69
 
     def test_precession_simple(
         self, diffraction_calculator_precession_simple, local_structure
@@ -148,7 +165,7 @@ def test_precession_simple(
             local_structure,
             reciprocal_radius=5.0,
         )
-        assert diffraction.coordinates.size == 250
+        assert diffraction.coordinates.size == 249
 
     def test_precession_full(
         self, diffraction_calculator_precession_full, local_structure
@@ -157,13 +174,15 @@ def test_precession_full(
             local_structure,
             reciprocal_radius=5.0,
         )
-        assert diffraction.coordinates.size == 250
+        assert diffraction.coordinates.size == 249
 
     def test_custom_shape_func(self, diffraction_calculator_custom, local_structure):
         diffraction = diffraction_calculator_custom.calculate_diffraction2d(
             local_structure,
             reciprocal_radius=5.0,
         )
+        # This shape factor model yields 0 intensity for the direct beam,
+        # which is less than the intensity threshold. The direct beam is therefore removed
         assert diffraction.coordinates.size == 52
 
     def test_appropriate_scaling(self, diffraction_calculator: SimulationGenerator):

From 909964cb2f5b31a8c95a63f5658536fa9553dac4 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Fri, 13 Dec 2024 10:54:12 +0100
Subject: [PATCH 08/17] Use rotate_with_basis

---
 diffsims/generators/simulation_generator.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index d0420708..58abedec 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -24,7 +24,7 @@
 from orix.quaternion import Rotation
 from orix.crystal_map import Phase
 
-from diffsims.crystallography._diffracting_vector import DiffractingVector, RotatedPhase
+from diffsims.crystallography._diffracting_vector import DiffractingVector
 from diffsims.utils.shape_factor_models import (
     linear,
     atanc,
@@ -437,11 +437,7 @@ def get_d(_c):
                 d > (lower - max_excitation_error)
             )
         # select these reflections
-        intersected_vectors = ~rot * (recip[intersection].to_miller())
-        intersected_vectors = DiffractingVector(
-            phase=RotatedPhase(recip.phase, rot),
-            xyz=intersected_vectors.data,
-        )
+        intersected_vectors = recip[intersection].rotate_with_basis(rot)
         excitation_error = excitation_error[intersection]
         hkl = recip_hkl[intersection]
         r_spot = intersected_vectors.norm

From 607cc449780fec800b9d9be7f0afadd4af825a9a Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Fri, 13 Dec 2024 13:47:17 +0100
Subject: [PATCH 09/17] Return self.__class__ instead of explicit
 ReciprocalLatticeVector

---
 diffsims/crystallography/_diffracting_vector.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/diffsims/crystallography/_diffracting_vector.py b/diffsims/crystallography/_diffracting_vector.py
index 46e6c7ea..27510dc5 100644
--- a/diffsims/crystallography/_diffracting_vector.py
+++ b/diffsims/crystallography/_diffracting_vector.py
@@ -183,7 +183,7 @@ def rotate_with_basis(self, rotation):
 
         # rotate vectors
         vecs = ~rotation * self.to_miller()
-        return ReciprocalLatticeVector(new_phase, xyz=vecs.data)
+        return self.__class__(new_phase, xyz=vecs.data)
 
     @property
     def intensity(self):

From 16509b7daf4318c9c27477e280e59d9825e1f6b4 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Thu, 23 Jan 2025 17:18:58 +0100
Subject: [PATCH 10/17] Adress #234

---
 .../crystallography/_diffracting_vector.py    |  10 +-
 .../reciprocal_lattice_vector.py              |  13 +-
 diffsims/generators/simulation_generator.py   | 297 +++++++++++-------
 .../generators/test_simulation_generator.py   |   2 +-
 4 files changed, 201 insertions(+), 121 deletions(-)

diff --git a/diffsims/crystallography/_diffracting_vector.py b/diffsims/crystallography/_diffracting_vector.py
index 27510dc5..5a516b88 100644
--- a/diffsims/crystallography/_diffracting_vector.py
+++ b/diffsims/crystallography/_diffracting_vector.py
@@ -202,11 +202,11 @@ def intensity(self, value):
             raise ValueError("Length of intensity array must match number of vectors")
         self._intensity = np.array(value)
 
-    def calculate_structure_factor(self):
-        raise NotImplementedError(
-            "Structure factor calculation not implemented for DiffractionVector. "
-            "Use ReciprocalLatticeVector instead."
-        )
+    # def calculate_structure_factor(self):
+    #     raise NotImplementedError(
+    #         "Structure factor calculation not implemented for DiffractionVector. "
+    #         "Use ReciprocalLatticeVector instead."
+    #     )
 
     def to_flat_polar(self):
         """Return the vectors in polar coordinates as projected onto the x,y plane"""
diff --git a/diffsims/crystallography/reciprocal_lattice_vector.py b/diffsims/crystallography/reciprocal_lattice_vector.py
index 23763f38..4fe22024 100644
--- a/diffsims/crystallography/reciprocal_lattice_vector.py
+++ b/diffsims/crystallography/reciprocal_lattice_vector.py
@@ -644,7 +644,9 @@ def allowed(self):
 
     # ------------------------- Custom methods ----------------------- #
 
-    def calculate_structure_factor(self, scattering_params="xtables"):
+    def calculate_structure_factor(
+        self, scattering_params="xtables", debye_waller_factors=None
+    ):
         r"""Populate :attr:`structure_factor` with the complex
         kinematical structure factor :math:`F_{hkl}` for each vector.
 
@@ -653,6 +655,8 @@ def calculate_structure_factor(self, scattering_params="xtables"):
         scattering_params : str
             Which atomic scattering factors to use, either ``"xtables"``
             (default) or ``"lobato"``.
+        debye_waller_factors: dict
+            Debye-Waller factors for atoms in the structure.
 
         Examples
         --------
@@ -712,6 +716,7 @@ def calculate_structure_factor(self, scattering_params="xtables"):
             structure=self.phase.structure,
             g_indices=hkl_unique,
             g_hkls_array=self.phase.structure.lattice.rnorm(hkl_unique),
+            debye_waller_factors=debye_waller_factors,
             scattering_params=scattering_params,
         )
 
@@ -933,6 +938,12 @@ def sanitise_phase(self):
         space_group = self.phase.space_group
         structure = self.phase.structure
 
+        # Lattice basis transform for hexagonal crystals can sometimes move atoms below 0
+        for atom in structure:
+            for i in range(3):
+                if -1e-6 < atom.xyz[i] < 0:
+                    atom.xyz[i] = 0
+
         new_structure = _expand_unit_cell(space_group, structure)
         for atom in new_structure:
             if np.issubdtype(type(atom.element), np.integer):
diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index 58abedec..d379c359 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -18,7 +18,7 @@
 
 """Kinematic Diffraction Simulation Generator."""
 
-from typing import Union, Sequence
+from typing import Union, Sequence, Tuple
 import numpy as np
 
 from orix.quaternion import Rotation
@@ -195,34 +195,44 @@ def calculate_diffraction2d(
                 min_dspacing=1 / reciprocal_radius,
                 include_zero_vector=with_direct_beam,
             )
+            recip.sanitise_phase()
+            # No need to pre-calculate structure factors if only a few rotations are present
+            if rotate.size > 4:
+                recip.calculate_structure_factor(
+                    scattering_params=self.scattering_params,
+                    debye_waller_factors=debye_waller_factors,
+                    )
             phase_vectors = []
-            # Avoid re-calculating in get_intersecting_reflections
-            recip_hkl = recip.hkl
             for rot in rotate:
                 # Calculate the reciprocal lattice vectors that intersect the Ewald sphere.
-                (
-                    intersected_vectors,
-                    hkl,
-                    shape_factor,
-                ) = self.get_intersecting_reflections(
+                intersection, excitation_error = get_intersection_with_ewalds_sphere(
                     recip,
-                    recip_hkl,
                     rot,
                     wavelength,
                     max_excitation_error,
-                    shape_factor_width=shape_factor_width,
-                    with_direct_beam=False,  # Already handled in DiffractingVector.from_min_dspacing
+                    self.precession_angle,
                 )
-
-                # Calculate diffracted intensities based on a kinematic model.
-                intensities = get_kinematical_intensities(
-                    p.structure,
-                    hkl,
-                    intersected_vectors.gspacing,
-                    prefactor=shape_factor,
-                    scattering_params=self.scattering_params,
-                    debye_waller_factors=debye_waller_factors,
+                # Select intersected reflections
+                intersected_vectors = recip[intersection].rotate_with_basis(rot)
+                excitation_error = excitation_error[intersection]
+                r_spot = intersected_vectors.norm
+
+                # Calculate shape factor
+                shape_factor = self.get_shape_factor(
+                    excitation_error, max_excitation_error, r_spot, shape_factor_width
                 )
+                # Calculate intensity
+                if rotate.size <= 4:
+                    v = recip[intersection]
+                    v.calculate_structure_factor(
+                        scattering_params=self.scattering_params,
+                        debye_waller_factors=debye_waller_factors,
+                        )
+                    f_hkls = v.structure_factor
+                else:
+                    f_hkls = recip[intersection].structure_factor
+                intensities = (f_hkls * f_hkls.conjugate()).real * shape_factor
+
                 # Threshold peaks included in simulation as factor of zero beam intensity.
                 peak_mask = intensities > np.max(intensities) * self.minimum_intensity
                 intensities = intensities[peak_mask]
@@ -334,6 +344,45 @@ def calculate_diffraction1d(
             wavelength=self.wavelength,
         )
 
+    def get_shape_factor(
+        self,
+        excitation_error: np.ndarray,
+        max_excitation_error: float,
+        r_spot: np.ndarray = None,
+        shape_factor_width: float = None,
+    ) -> np.ndarray:
+
+        if shape_factor_width is None:
+            shape_factor_width = max_excitation_error
+        # select and evaluate shape factor model
+        if self.precession_angle == 0:
+            # calculate shape factor
+            shape_factor = self.shape_factor_model(
+                excitation_error, shape_factor_width, **self.shape_factor_kwargs
+            )
+        else:
+            if self.approximate_precession:
+                shape_factor = lorentzian_precession(
+                    excitation_error,
+                    shape_factor_width,
+                    r_spot,
+                    np.deg2rad(self.precession_angle),
+                )
+            else:
+                if r_spot is None:
+                    raise ValueError(
+                        "Must supply `r_spot` parameter when not approximating precession"
+                    )
+                shape_factor = _shape_factor_precession(
+                    excitation_error,
+                    r_spot,
+                    np.deg2rad(self.precession_angle),
+                    self.shape_factor_model,
+                    shape_factor_width,
+                    **self.shape_factor_kwargs,
+                )
+        return shape_factor
+
     def get_intersecting_reflections(
         self,
         recip: DiffractingVector,
@@ -363,108 +412,128 @@ 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.
         """
-        # Identify the excitation errors of all points (distance from point to Ewald sphere)
-
-        # Instead of rotating vectors, rotate Ewald's sphere to find intersections.
-        # Only rotate the intersecting vectors.
-        # Using notation from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
-        r = 1 / wavelength
-        u = rot.to_matrix().squeeze() @ np.array([0, 0, 1])
-        c = r * u
-        o = recip.data
         if with_direct_beam:
-            o = np.vstack([o, [0, 0, 0]])
             # Modify the data directly, as `data` is a property
             recip._data = np.vstack([recip._data, [0, 0, 0]])
             recip_hkl = np.vstack([recip_hkl, [0, 0, 0]])
 
-        diff = o - c
-        dot = np.dot(u, diff.T)
-        nabla = dot**2 - np.sum(diff**2, axis=1) + r**2
-        # We know the reflections are all going to intersect twice
-        # Therefore, no need to look at the sign of nabla
-        # We also know only the smaller root is important, i.e. -sqrt(nabla)
-        sqrt_nabla = np.sqrt(nabla)
-        d = -dot - sqrt_nabla
-
-        excitation_error = d
-
-        # determine the pre-selection reflections
-        if self.precession_angle == 0:
-            intersection = np.abs(excitation_error) < max_excitation_error
-        else:
-            # Using the following script to find equations for upper and lower bounds for precessing Ewald's sphere
-            """
-            import sympy
-            import numpy as np
-
-            a = sympy.Symbol("a") # Precession angle
-            r = sympy.Symbol("r") # Ewald's sphere radius
-            rho, z = sympy.symbols("rho z") # cylindrical coordinates of reflection
-
-            rot = lambda ang: np.asarray([[sympy.cos(ang), -sympy.sin(ang)],[sympy.sin(ang), sympy.cos(ang)]])
-
-            u = np.asarray([0, 1])
-            c = r * u
-            cl = rot(a) @ c
-            cr = rot(-a) @ c
-            o = np.asarray([rho, z])
-
-            def get_d(_c):
-                diff = o - _c
-                dot = np.dot(u, diff)
-                nabla = dot**2 - sum(i**2 for i in diff) + r**2
-                sqrt_nabla = nabla**0.5
-                return  -dot - sqrt_nabla
-
-            d = get_d(c)
-            d_upper = get_d(cl)
-            d_lower = get_d(cr)
-
-            print(d.simplify())             # r - z - (r**2 - rho**2)**0.5
-            print((d_upper - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) + rho)**2)**0.5
-            print((d_lower - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) - rho)**2)**0.5
-            """
-            # In the above script, d is the same as before.
-            # We need the distance of the reflections from the incident beam, i.e. the cylindrical coordinate rho
-            # (using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation):
-            rho = np.linalg.norm(np.dot(o, u)[:, np.newaxis] * u - o, axis=1)
-            a = np.deg2rad(self.precession_angle)
-            first_half = r * np.cos(a) - r + (r**2 - rho**2) ** 0.5
-            upper = first_half - (r**2 - (r * np.sin(a) + rho) ** 2) ** 0.5
-            lower = first_half - (r**2 - (r * np.sin(a) - rho) ** 2) ** 0.5
-            intersection = (d < (upper + max_excitation_error)) & (
-                d > (lower - max_excitation_error)
-            )
-        # select these reflections
+        intersection, excitation_error = get_intersection_with_ewalds_sphere(
+            recip,
+            rot,
+            wavelength,
+            max_excitation_error,
+            self.precession_angle,
+        )
         intersected_vectors = recip[intersection].rotate_with_basis(rot)
         excitation_error = excitation_error[intersection]
         hkl = recip_hkl[intersection]
         r_spot = intersected_vectors.norm
 
-        if shape_factor_width is None:
-            shape_factor_width = max_excitation_error
-        # select and evaluate shape factor model
-        if self.precession_angle == 0:
-            # calculate shape factor
-            shape_factor = self.shape_factor_model(
-                excitation_error, shape_factor_width, **self.shape_factor_kwargs
-            )
-        else:
-            if self.approximate_precession:
-                shape_factor = lorentzian_precession(
-                    excitation_error,
-                    shape_factor_width,
-                    r_spot,
-                    np.deg2rad(self.precession_angle),
-                )
-            else:
-                shape_factor = _shape_factor_precession(
-                    excitation_error,
-                    r_spot,
-                    np.deg2rad(self.precession_angle),
-                    self.shape_factor_model,
-                    shape_factor_width,
-                    **self.shape_factor_kwargs,
-                )
+        shape_factor = self.get_shape_factor(
+            excitation_error, max_excitation_error, r_spot, shape_factor_width
+        )
         return intersected_vectors, hkl, shape_factor
+
+
+# TODO consider refactoring into a seperate file
+def get_intersection_with_ewalds_sphere(
+    recip: DiffractingVector,
+    rot: Rotation,
+    wavelength: float,
+    max_excitation_error: float,
+    precession_angle: float = 0,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Calculates the reciprocal lattice vectors that intersect the Ewald sphere.
+
+    Parameters
+    ----------
+    recip
+        The reciprocal lattice vectors to rotate.
+    rot
+        The rotation to apply to the reciprocal lattice vectors.
+    wavelength
+        The wavelength of the electrons in Angstroms.
+    max_excitation_error
+        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.
+    precession_angle
+        Degrees
+
+    Returns
+    -------
+    intersection
+        Array of bools. True where the vectors intersect
+    excitation_error
+        Excitation error of all vectors
+    """
+    # Identify the excitation errors of all points (distance from point to Ewald sphere)
+
+    # Instead of rotating vectors, rotate Ewald's sphere to find intersections.
+    # Only rotate the intersecting vectors.
+    # Using notation from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
+    r = 1 / wavelength
+    u = rot.to_matrix().squeeze() @ np.array([0, 0, 1])
+    c = r * u
+    o = recip.data
+
+    diff = o - c
+    dot = np.dot(u, diff.T)
+    nabla = dot**2 - np.sum(diff**2, axis=1) + r**2
+    # We know the reflections are all going to intersect twice
+    # Therefore, no need to look at the sign of nabla
+    # We also know only the smaller root is important, i.e. -sqrt(nabla)
+    sqrt_nabla = np.sqrt(nabla)
+    d = -dot - sqrt_nabla
+
+    excitation_error = d
+
+    # determine the pre-selection reflections
+    if precession_angle == 0:
+        intersection = np.abs(excitation_error) < max_excitation_error
+    else:
+        # Using the following script to find equations for upper and lower bounds for precessing Ewald's sphere
+        """
+        import sympy
+        import numpy as np
+
+        a = sympy.Symbol("a") # Precession angle
+        r = sympy.Symbol("r") # Ewald's sphere radius
+        rho, z = sympy.symbols("rho z") # cylindrical coordinates of reflection
+
+        rot = lambda ang: np.asarray([[sympy.cos(ang), -sympy.sin(ang)],[sympy.sin(ang), sympy.cos(ang)]])
+
+        u = np.asarray([0, 1])
+        c = r * u
+        cl = rot(a) @ c
+        cr = rot(-a) @ c
+        o = np.asarray([rho, z])
+
+        def get_d(_c):
+            diff = o - _c
+            dot = np.dot(u, diff)
+            nabla = dot**2 - sum(i**2 for i in diff) + r**2
+            sqrt_nabla = nabla**0.5
+            return  -dot - sqrt_nabla
+
+        d = get_d(c)
+        d_upper = get_d(cl)
+        d_lower = get_d(cr)
+
+        print(d.simplify())             # r - z - (r**2 - rho**2)**0.5
+        print((d_upper - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) + rho)**2)**0.5
+        print((d_lower - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) - rho)**2)**0.5
+        """
+        # In the above script, d is the same as before.
+        # We need the distance of the reflections from the incident beam, i.e. the cylindrical coordinate rho
+        # (using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation):
+        rho = np.linalg.norm(np.dot(o, u)[:, np.newaxis] * u - o, axis=1)
+        a = np.deg2rad(precession_angle)
+        first_half = r * np.cos(a) - r + (r**2 - rho**2) ** 0.5
+        upper = first_half - (r**2 - (r * np.sin(a) + rho) ** 2) ** 0.5
+        lower = first_half - (r**2 - (r * np.sin(a) - rho) ** 2) ** 0.5
+        intersection = (d < (upper + max_excitation_error)) & (
+            d > (lower - max_excitation_error)
+        )
+    return intersection, excitation_error
diff --git a/diffsims/tests/generators/test_simulation_generator.py b/diffsims/tests/generators/test_simulation_generator.py
index 67f54723..b6036388 100644
--- a/diffsims/tests/generators/test_simulation_generator.py
+++ b/diffsims/tests/generators/test_simulation_generator.py
@@ -345,7 +345,7 @@ def test_same_simulation_results():
     # old_data = diff_lib["Graphite"]["simulations"][0].get_diffraction_pattern(shape=shape, sigma=sigma)
 
     # New
-    p = Phase("Graphite", point_group="6/mmm", structure=structure_matrix)
+    p = Phase("Graphite", structure=structure_matrix, space_group=1) # Space group 1 ensures no symmetry is applied
     gen = SimulationGenerator(**generator_kwargs)
     rot = Rotation.from_euler(euler_angles_new, degrees=True)
     sim = gen.calculate_diffraction2d(

From 8a7b31421b31b29fd5903027595f1c2345cc6db7 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Fri, 24 Jan 2025 16:49:56 +0100
Subject: [PATCH 11/17] Getting closer, still some tests left

---
 .../reciprocal_lattice_vector.py              | 77 +++++++++----------
 diffsims/generators/simulation_generator.py   | 24 ++----
 .../generators/test_simulation_generator.py   | 15 ++--
 3 files changed, 52 insertions(+), 64 deletions(-)

diff --git a/diffsims/crystallography/reciprocal_lattice_vector.py b/diffsims/crystallography/reciprocal_lattice_vector.py
index 4fe22024..4893d621 100644
--- a/diffsims/crystallography/reciprocal_lattice_vector.py
+++ b/diffsims/crystallography/reciprocal_lattice_vector.py
@@ -20,7 +20,7 @@
 from copy import deepcopy
 
 from diffpy.structure.symmetryutilities import expandPosition
-from diffpy.structure import Structure
+from diffpy.structure import Structure, Lattice
 import numba as nb
 import numpy as np
 from orix.vector import Miller, Vector3d
@@ -696,22 +696,8 @@ def calculate_structure_factor(
         """
 
         # Compute one structure factor per set {hkl}
-        hkl_sets = self.get_hkl_sets()
-
-        # For each set, get the indices of the first vector in the
-        # present vectors, accounting for potential multiple dimensions
-        # and avoding computing the unique vectors again
-        first_idx = []
-        for arr in list(hkl_sets.values()):
-            i = []
-            for arr_i in arr:
-                i.append(arr_i[0])
-            first_idx.append(i)
-        first_idx_arr = np.array(first_idx).T
-
-        # Get 2D array of unique vectors, one for each set
-        hkl_unique = self.hkl[tuple(first_idx_arr)]
-
+        unique, inds = self.unique(use_symmetry=True, return_inverse=True)
+        hkl_unique = unique.hkl
         structure_factor = _get_kinematical_structure_factor(
             structure=self.phase.structure,
             g_indices=hkl_unique,
@@ -719,10 +705,8 @@ def calculate_structure_factor(
             debye_waller_factors=debye_waller_factors,
             scattering_params=scattering_params,
         )
-
         # Set structure factors of all symmetrically equivalent vectors
-        for i, idx in enumerate(hkl_sets.values()):
-            self._structure_factor[idx] = structure_factor[i]
+        self._structure_factor[:] = structure_factor[inds]
 
     def calculate_theta(self, voltage):
         r"""Populate :attr:`theta` with the Bragg angle :math:`theta_B`
@@ -1517,7 +1501,7 @@ def transpose(self, *axes):
         new._update_shapes()
         return new
 
-    def unique(self, use_symmetry=False, return_index=False):
+    def unique(self, use_symmetry=False, return_index=False, return_inverse=False):
         """The unique vectors.
 
         Parameters
@@ -1528,6 +1512,9 @@ def unique(self, use_symmetry=False, return_index=False):
         return_index : bool, optional
             Whether to return the indices of the (flattened) data where
             the unique entries were found. Default is ``False``.
+        return_inverse: bool, optional
+            Whether to also return the indices to reconstruct the
+            (flattened) data from the unique data.
 
         Returns
         -------
@@ -1535,41 +1522,41 @@ def unique(self, use_symmetry=False, return_index=False):
             Flattened unique vectors.
         idx : numpy.ndarray
             Indices of the unique data in the (flattened) array.
-
+        inv : numpy.ndarray
+            Indices to reconstruct the original vectors from the unique
         """
 
         # TODO: Reduce floating point precision in orix!
         def miller_unique(miller, use_symmetry=False):
-            v, idx = Vector3d(miller).unique(return_index=True)
+            v, idx, inv = Vector3d(miller).unique(return_index=True, return_inverse=True)
+            idx = idx[::-1]
+            inv = inv[::-1]
 
             if use_symmetry:
                 operations = miller.phase.point_group
                 n_v = v.size
-                v2 = operations.outer(v).flatten().reshape(*(n_v, operations.size))
+                v2 = operations.outer(v[::-1]).flatten().reshape(*(n_v, operations.size))
                 data = v2.data.round(6)  # Reduced precision
-                data_sorted = np.zeros_like(data)
-                for i in range(n_v):
-                    a = data[i]
-                    order = np.lexsort(a.T)  # Sort by column 1, 2, then 3
-                    data_sorted[i] = a[order]
-                _, idx = np.unique(data_sorted, return_index=True, axis=0)
-                v = v[idx[::-1]]
+                _, idx, inv = np.unique(data, return_index=True, return_inverse=True, axis=0)
+                v = v[idx]
 
             m = miller.__class__(xyz=v.data, phase=miller.phase)
             m.coordinate_format = miller.coordinate_format
-            return m, idx
 
-        #        kwargs = dict(use_symmetry=use_symmetry, return_index=True)
-        #        miller, idx = self.to_miller().unique(**kwargs)
-        miller, idx = miller_unique(self.to_miller(), use_symmetry)
-        idx = idx[::-1]
+            return m, idx, inv
+
+        miller, idx, inv = miller_unique(self.to_miller(), use_symmetry)
 
         new_rlv = self.from_miller(miller)
         new_rlv._structure_factor = self.structure_factor.ravel()[idx]
         new_rlv._theta = self.theta.ravel()[idx]
 
-        if return_index:
+        if return_index and return_inverse:
+            return new_rlv, idx, inv
+        elif return_index:
             return new_rlv, idx
+        elif return_inverse:
+            return new_rlv, inv
         else:
             return new_rlv
 
@@ -1623,7 +1610,7 @@ def _atom_eq(atom1, atom2):
 
     return (
         atom1.element == atom2.element
-        and np.allclose(atom1.xyz, atom2.xyz, atol=1e-7)
+        and np.allclose(atom1.xyz_cartn, atom2.xyz_cartn, atol=1e-7)
         and atom1.occupancy == atom2.occupancy
         and np.allclose(atom1.U, atom2.U, atol=1e-7)
         and np.allclose(atom1.Uisoequiv, atom2.Uisoequiv, atol=1e-7)
@@ -1646,9 +1633,17 @@ def _expand_unit_cell(space_group, structure):
     new_structure : diffpy.structure.Structure
 
     """
-
+    # Transform to diffpy axis conventions
+    structure_matrix = structure.lattice.base
+    pos = structure.xyz_cartn
+    structure = Structure(
+        atoms=list(structure),
+        lattice=Lattice(*structure.lattice.abcABG()),
+    )
     new_structure = Structure(lattice=structure.lattice)
+    structure.xyz_cartn = pos
 
+    # Perform symmetry expansion
     for atom in structure:
         equal = []
         for atom2 in new_structure:
@@ -1660,6 +1655,10 @@ def _expand_unit_cell(space_group, structure):
                 new_atom.xyz = new_position
                 new_structure.append(new_atom)
 
+    # Transform back to orix convention
+    old_xyz_cartn = new_structure.xyz_cartn
+    new_structure.lattice.setLatBase(structure_matrix)
+    new_structure.xyz_cartn = old_xyz_cartn
     return new_structure
 
 
diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index d379c359..96ae4728 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -131,9 +131,7 @@ def wavelength(self):
     def calculate_diffraction2d(
         self,
         phase: Union[Phase, Sequence[Phase]],
-        rotation: Union[Rotation, Sequence[Rotation]] = Rotation.from_euler(
-            (0, 0, 0), degrees=True
-        ),
+        rotation: Union[Rotation, Sequence[Rotation]] = Rotation.identity(),
         reciprocal_radius: float = 1.0,
         with_direct_beam: bool = True,
         max_excitation_error: float = 1e-2,
@@ -196,12 +194,10 @@ def calculate_diffraction2d(
                 include_zero_vector=with_direct_beam,
             )
             recip.sanitise_phase()
-            # No need to pre-calculate structure factors if only a few rotations are present
-            if rotate.size > 4:
-                recip.calculate_structure_factor(
-                    scattering_params=self.scattering_params,
-                    debye_waller_factors=debye_waller_factors,
-                    )
+            recip.calculate_structure_factor(
+                scattering_params=self.scattering_params,
+                debye_waller_factors=debye_waller_factors,
+            )
             phase_vectors = []
             for rot in rotate:
                 # Calculate the reciprocal lattice vectors that intersect the Ewald sphere.
@@ -222,15 +218,7 @@ def calculate_diffraction2d(
                     excitation_error, max_excitation_error, r_spot, shape_factor_width
                 )
                 # Calculate intensity
-                if rotate.size <= 4:
-                    v = recip[intersection]
-                    v.calculate_structure_factor(
-                        scattering_params=self.scattering_params,
-                        debye_waller_factors=debye_waller_factors,
-                        )
-                    f_hkls = v.structure_factor
-                else:
-                    f_hkls = recip[intersection].structure_factor
+                f_hkls = recip.structure_factor[intersection]
                 intensities = (f_hkls * f_hkls.conjugate()).real * shape_factor
 
                 # Threshold peaks included in simulation as factor of zero beam intensity.
diff --git a/diffsims/tests/generators/test_simulation_generator.py b/diffsims/tests/generators/test_simulation_generator.py
index b6036388..d1538d47 100644
--- a/diffsims/tests/generators/test_simulation_generator.py
+++ b/diffsims/tests/generators/test_simulation_generator.py
@@ -195,12 +195,13 @@ def test_appropriate_scaling(self, diffraction_calculator: SimulationGenerator):
         big_diffraction = diffraction_calculator.calculate_diffraction2d(
             phase=big_silicon, reciprocal_radius=5.0
         )
-        indices = [tuple(i) for i in diffraction.coordinates.hkl]
-        big_indices = [tuple(i) for i in big_diffraction.coordinates.hkl]
-        assert (2, 2, 0) in indices
-        assert (2, 2, 0) in big_indices
-        coordinates = diffraction.coordinates[indices.index((2, 2, 0))]
-        big_coordinates = big_diffraction.coordinates[big_indices.index((2, 2, 0))]
+        indices = [tuple(i.tolist()) for i in diffraction.coordinates.hkl.astype(int)]
+        big_indices = [tuple(i.tolist()) for i in big_diffraction.coordinates.hkl.astype(int)]
+        target = (2, 2, 0)
+        assert target in indices
+        assert target in big_indices
+        coordinates = diffraction.coordinates[indices.index(target)]
+        big_coordinates = big_diffraction.coordinates[big_indices.index(target)]
         assert np.allclose(coordinates.data, big_coordinates.data * 2)
 
     def test_appropriate_intensities(self, diffraction_calculator, local_structure):
@@ -345,7 +346,7 @@ def test_same_simulation_results():
     # old_data = diff_lib["Graphite"]["simulations"][0].get_diffraction_pattern(shape=shape, sigma=sigma)
 
     # New
-    p = Phase("Graphite", structure=structure_matrix, space_group=1) # Space group 1 ensures no symmetry is applied
+    p = Phase("Graphite", structure=structure_matrix, space_group=194)
     gen = SimulationGenerator(**generator_kwargs)
     rot = Rotation.from_euler(euler_angles_new, degrees=True)
     sim = gen.calculate_diffraction2d(

From 6ef265ff37074b64f426bf17ccaece9c61954dc4 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Sat, 25 Jan 2025 09:36:19 +0100
Subject: [PATCH 12/17] Ensure compatible ordering of unique vectors

---
 .../crystallography/reciprocal_lattice_vector.py   | 14 +++++++++-----
 .../crystallography/test_diffracting_vector.py     |  4 ++--
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/diffsims/crystallography/reciprocal_lattice_vector.py b/diffsims/crystallography/reciprocal_lattice_vector.py
index 4893d621..b7eeb117 100644
--- a/diffsims/crystallography/reciprocal_lattice_vector.py
+++ b/diffsims/crystallography/reciprocal_lattice_vector.py
@@ -1529,15 +1529,19 @@ def unique(self, use_symmetry=False, return_index=False, return_inverse=False):
         # TODO: Reduce floating point precision in orix!
         def miller_unique(miller, use_symmetry=False):
             v, idx, inv = Vector3d(miller).unique(return_index=True, return_inverse=True)
-            idx = idx[::-1]
-            inv = inv[::-1]
 
             if use_symmetry:
                 operations = miller.phase.point_group
                 n_v = v.size
-                v2 = operations.outer(v[::-1]).flatten().reshape(*(n_v, operations.size))
-                data = v2.data.round(6)  # Reduced precision
-                _, idx, inv = np.unique(data, return_index=True, return_inverse=True, axis=0)
+                v2 = operations.outer(v).flatten().reshape(*(n_v, operations.size))
+                data = v2.data.round(8)  # Reduced precision
+                # To check for uniqueness, sort the equivalent vectors from each operation
+                data_sorted = np.zeros_like(data)
+                for i in range(n_v):
+                    a = data[i]
+                    order = np.lexsort(a.T)  # Sort by column 1, 2, then 3
+                    data_sorted[i] = a[order[::-1]] # Reverse to preserve order
+                _, idx, inv = np.unique(data_sorted, return_index=True, return_inverse=True, axis=0)
                 v = v[idx]
 
             m = miller.__class__(xyz=v.data, phase=miller.phase)
diff --git a/diffsims/tests/crystallography/test_diffracting_vector.py b/diffsims/tests/crystallography/test_diffracting_vector.py
index bf4ccdf6..5a234c34 100644
--- a/diffsims/tests/crystallography/test_diffracting_vector.py
+++ b/diffsims/tests/crystallography/test_diffracting_vector.py
@@ -41,8 +41,8 @@ def test_slicing(self, ferrite_phase):
 
     def test_structure_factor(self, ferrite_phase):
         rlv = DiffractingVector.from_min_dspacing(ferrite_phase, 1.5)
-        with pytest.raises(NotImplementedError):
-            rlv.calculate_structure_factor()
+        rlv.calculate_structure_factor()
+        assert np.all(~np.isnan(rlv.structure_factor))
 
     def test_hkl(self, ferrite_phase):
         rlv = ReciprocalLatticeVector(ferrite_phase, hkl=[[1, 1, 1], [2, 0, 0]])

From 54b77f64d9b788df0184f17786971a193e994068 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Sat, 25 Jan 2025 12:44:48 +0100
Subject: [PATCH 13/17] Fix test by using symmetrically unique reflections (4
 degenerate were missing)

---
 diffsims/tests/generators/test_simulation_generator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/diffsims/tests/generators/test_simulation_generator.py b/diffsims/tests/generators/test_simulation_generator.py
index d1538d47..8afdb1da 100644
--- a/diffsims/tests/generators/test_simulation_generator.py
+++ b/diffsims/tests/generators/test_simulation_generator.py
@@ -183,7 +183,7 @@ def test_custom_shape_func(self, diffraction_calculator_custom, local_structure)
         )
         # This shape factor model yields 0 intensity for the direct beam,
         # which is less than the intensity threshold. The direct beam is therefore removed
-        assert diffraction.coordinates.size == 52
+        assert diffraction.coordinates.unique(use_symmetry=True).size == 9
 
     def test_appropriate_scaling(self, diffraction_calculator: SimulationGenerator):
         """Tests that doubling the unit cell halves the pattern spacing."""

From a55784dfda5aa0b8a5fc8871d38d33d605898401 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Mon, 27 Jan 2025 14:08:59 +0100
Subject: [PATCH 14/17] Use Phase.expand_asymmetric_unit

---
 .../crystallography/_diffracting_vector.py    |  3 +
 .../reciprocal_lattice_vector.py              | 94 ++-----------------
 .../generators/test_simulation_generator.py   | 44 ++++++++-
 3 files changed, 55 insertions(+), 86 deletions(-)

diff --git a/diffsims/crystallography/_diffracting_vector.py b/diffsims/crystallography/_diffracting_vector.py
index 5a516b88..0b8b17e9 100644
--- a/diffsims/crystallography/_diffracting_vector.py
+++ b/diffsims/crystallography/_diffracting_vector.py
@@ -37,7 +37,10 @@ class RotatedPhase(Phase):
     """
 
     def __init__(self, phase: Phase, rotation: Rotation):
+        # Copy structure to not override the input phase.
+        # Use private members to avoid re-computing
         self._structure = Structure(phase.structure)
+        self._diffpy_lattice = phase._diffpy_lattice
         self.name = phase.name
         self.space_group = phase.space_group
         self.point_group = phase.point_group
diff --git a/diffsims/crystallography/reciprocal_lattice_vector.py b/diffsims/crystallography/reciprocal_lattice_vector.py
index b7eeb117..75e66cf0 100644
--- a/diffsims/crystallography/reciprocal_lattice_vector.py
+++ b/diffsims/crystallography/reciprocal_lattice_vector.py
@@ -19,8 +19,6 @@
 from collections import defaultdict
 from copy import deepcopy
 
-from diffpy.structure.symmetryutilities import expandPosition
-from diffpy.structure import Structure, Lattice
 import numba as nb
 import numpy as np
 from orix.vector import Miller, Vector3d
@@ -919,22 +917,12 @@ def sanitise_phase(self):
 
         self._raise_if_no_space_group()
 
-        space_group = self.phase.space_group
-        structure = self.phase.structure
+        self.phase.expand_asymmetric_unit()
 
-        # Lattice basis transform for hexagonal crystals can sometimes move atoms below 0
-        for atom in structure:
-            for i in range(3):
-                if -1e-6 < atom.xyz[i] < 0:
-                    atom.xyz[i] = 0
-
-        new_structure = _expand_unit_cell(space_group, structure)
-        for atom in new_structure:
+        for atom in self.phase.structure:
             if np.issubdtype(type(atom.element), np.integer):
                 atom.element = _get_string_from_element_id(atom.element)
 
-        self.phase.structure = new_structure
-
     def symmetrise(self, return_multiplicity=False, return_index=False):
         """Unique vectors symmetrically equivalent to the vectors.
 
@@ -1528,7 +1516,9 @@ def unique(self, use_symmetry=False, return_index=False, return_inverse=False):
 
         # TODO: Reduce floating point precision in orix!
         def miller_unique(miller, use_symmetry=False):
-            v, idx, inv = Vector3d(miller).unique(return_index=True, return_inverse=True)
+            v, idx, inv = Vector3d(miller).unique(
+                return_index=True, return_inverse=True, ignore_zero=False
+            )
 
             if use_symmetry:
                 operations = miller.phase.point_group
@@ -1540,8 +1530,10 @@ def miller_unique(miller, use_symmetry=False):
                 for i in range(n_v):
                     a = data[i]
                     order = np.lexsort(a.T)  # Sort by column 1, 2, then 3
-                    data_sorted[i] = a[order[::-1]] # Reverse to preserve order
-                _, idx, inv = np.unique(data_sorted, return_index=True, return_inverse=True, axis=0)
+                    data_sorted[i] = a[order[::-1]]  # Reverse to preserve order
+                _, idx, inv = np.unique(
+                    data_sorted, return_index=True, return_inverse=True, axis=0
+                )
                 v = v[idx]
 
             m = miller.__class__(xyz=v.data, phase=miller.phase)
@@ -1598,74 +1590,6 @@ def stack(cls, sequence):
         return new
 
 
-# TODO: Upstream to diffpy.structure.Atom.__eq__()
-def _atom_eq(atom1, atom2):
-    """Determine whether two atoms are equal.
-
-    Parameters
-    ----------
-    atom1, atom2 : diffpy.structure.Atom
-
-    Returns
-    -------
-    bool
-
-    """
-
-    return (
-        atom1.element == atom2.element
-        and np.allclose(atom1.xyz_cartn, atom2.xyz_cartn, atol=1e-7)
-        and atom1.occupancy == atom2.occupancy
-        and np.allclose(atom1.U, atom2.U, atol=1e-7)
-        and np.allclose(atom1.Uisoequiv, atom2.Uisoequiv, atol=1e-7)
-    )
-
-
-# TODO: Upstream to orix.crystal_map.Phase.expand_structure()
-def _expand_unit_cell(space_group, structure):
-    """Expand a unit cell with symmetrically equivalent atoms.
-
-    Parameters
-    ----------
-    space_group : diffpy.structure.spacegroupmod.SpaceGroup
-        Space group describing the symmetry operations of the unit cell.
-    structure : diffpy.structure.Structure
-        Initial structure with atoms.
-
-    Returns
-    -------
-    new_structure : diffpy.structure.Structure
-
-    """
-    # Transform to diffpy axis conventions
-    structure_matrix = structure.lattice.base
-    pos = structure.xyz_cartn
-    structure = Structure(
-        atoms=list(structure),
-        lattice=Lattice(*structure.lattice.abcABG()),
-    )
-    new_structure = Structure(lattice=structure.lattice)
-    structure.xyz_cartn = pos
-
-    # Perform symmetry expansion
-    for atom in structure:
-        equal = []
-        for atom2 in new_structure:
-            equal.append(_atom_eq(atom, atom2))
-        if not any(equal):
-            new_positions = expandPosition(space_group, atom.xyz)[0]
-            for new_position in new_positions:
-                new_atom = deepcopy(atom)
-                new_atom.xyz = new_position
-                new_structure.append(new_atom)
-
-    # Transform back to orix convention
-    old_xyz_cartn = new_structure.xyz_cartn
-    new_structure.lattice.setLatBase(structure_matrix)
-    new_structure.xyz_cartn = old_xyz_cartn
-    return new_structure
-
-
 @nb.njit(
     "int64[:](float64[:, :], float64[:, :], int64[:])",
     cache=True,
diff --git a/diffsims/tests/generators/test_simulation_generator.py b/diffsims/tests/generators/test_simulation_generator.py
index 8afdb1da..4d861ff7 100644
--- a/diffsims/tests/generators/test_simulation_generator.py
+++ b/diffsims/tests/generators/test_simulation_generator.py
@@ -196,7 +196,9 @@ def test_appropriate_scaling(self, diffraction_calculator: SimulationGenerator):
             phase=big_silicon, reciprocal_radius=5.0
         )
         indices = [tuple(i.tolist()) for i in diffraction.coordinates.hkl.astype(int)]
-        big_indices = [tuple(i.tolist()) for i in big_diffraction.coordinates.hkl.astype(int)]
+        big_indices = [
+            tuple(i.tolist()) for i in big_diffraction.coordinates.hkl.astype(int)
+        ]
         target = (2, 2, 0)
         assert target in indices
         assert target in big_indices
@@ -361,3 +363,43 @@ def test_same_simulation_results():
     )
     old_data = np.load(FILE1)
     np.testing.assert_allclose(new_data, old_data, atol=1e-8)
+
+
+def test_space_group_is_accounted_for():
+    structure = diffpy.structure.Structure(
+        atoms=[diffpy.structure.Atom("Au", xyz=(0, 0, 0))],
+        lattice=diffpy.structure.Lattice(4, 4, 4, 90, 90, 90),
+    )
+    P1 = Phase(space_group=1, structure=structure)
+    Fd3m = Phase(space_group=227, structure=structure)
+
+    gen = SimulationGenerator()
+    P1_sim = gen.calculate_diffraction2d(P1)
+    Fd3m_sim = gen.calculate_diffraction2d(Fd3m)
+
+    # Use ReciprocalLatticeVector.sanitize_phase
+    Fd3m_expanded = Fd3m.deepcopy()
+    Fd3m_expanded.expand_asymmetric_unit()
+    Fd3m_expanded_sim = gen.calculate_diffraction2d(Fd3m_expanded)
+
+    # Check coordinates. There are many extinct reflections in Fd3m, so look for those in P1
+    P1_hkl = [tuple(hkl.round().astype(int).tolist()) for hkl in P1_sim.coordinates.hkl]
+    Fd3m_hkl = [
+        tuple(hkl.round().astype(int).tolist()) for hkl in Fd3m_sim.coordinates.hkl
+    ]
+    Fd3m_expanded_hkl = [
+        tuple(hkl.round().astype(int).tolist())
+        for hkl in Fd3m_expanded_sim.coordinates.hkl
+    ]
+    assert set(P1_hkl).issuperset(set(Fd3m_hkl))
+    assert set(P1_hkl).issuperset(set(Fd3m_expanded_hkl))
+    assert set(Fd3m_hkl) == set(Fd3m_expanded_hkl)
+
+    # Check intensities, should be different
+    order = [P1_hkl.index(hkl) for hkl in Fd3m_hkl]
+    assert not np.allclose(
+        P1_sim.coordinates.intensity[order], Fd3m_sim.coordinates.intensity
+    )
+    assert np.allclose(
+        Fd3m_sim.coordinates.intensity, Fd3m_expanded_sim.coordinates.intensity
+    )

From e3b09d854a99d2ae588b94f9305c6a783d9b2376 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Mon, 27 Jan 2025 16:03:35 +0100
Subject: [PATCH 15/17] Simplify `DiffractingVector.__getitem__`

---
 .../crystallography/_diffracting_vector.py    | 24 +++----------------
 1 file changed, 3 insertions(+), 21 deletions(-)

diff --git a/diffsims/crystallography/_diffracting_vector.py b/diffsims/crystallography/_diffracting_vector.py
index 0b8b17e9..85961645 100644
--- a/diffsims/crystallography/_diffracting_vector.py
+++ b/diffsims/crystallography/_diffracting_vector.py
@@ -124,27 +124,9 @@ def __init__(self, phase, xyz=None, hkl=None, hkil=None, intensity=None):
     def __getitem__(self, key):
         new_data = self.data[key]
         dv_new = self.__class__(self.phase, xyz=new_data)
-
-        if np.isnan(self.structure_factor).all():
-            dv_new._structure_factor = np.full(dv_new.shape, np.nan, dtype="complex128")
-
-        else:
-            dv_new._structure_factor = self.structure_factor[key]
-        if np.isnan(self.theta).all():
-            dv_new._theta = np.full(dv_new.shape, np.nan)
-        else:
-            dv_new._theta = self.theta[key]
-        if np.isnan(self.intensity).all():
-            dv_new._intensity = np.full(dv_new.shape, np.nan)
-        else:
-            slic = self.intensity[key]
-            if not hasattr(slic, "__len__"):
-                slic = np.array(
-                    [
-                        slic,
-                    ]
-                )
-            dv_new._intensity = slic
+        dv_new._structure_factor = self._structure_factor[key]
+        dv_new._theta = self._theta[key]
+        dv_new._intensity = self._intensity[key]
 
         return dv_new
 

From 7e307fb3281ef721d6758f66fa9a556a45adba6f Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Mon, 27 Jan 2025 16:04:51 +0100
Subject: [PATCH 16/17] Use numba for speedup

---
 diffsims/generators/simulation_generator.py | 172 +++++++++++++-------
 1 file changed, 114 insertions(+), 58 deletions(-)

diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index 96ae4728..ad5cb8df 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -20,8 +20,10 @@
 
 from typing import Union, Sequence, Tuple
 import numpy as np
+from numba import njit
 
 from orix.quaternion import Rotation
+from orix.vector import Vector3d
 from orix.crystal_map import Phase
 
 from diffsims.crystallography._diffracting_vector import DiffractingVector
@@ -199,11 +201,11 @@ def calculate_diffraction2d(
                 debye_waller_factors=debye_waller_factors,
             )
             phase_vectors = []
-            for rot in rotate:
+            for rot, optical_axis in zip(rotate, rotate * Vector3d.zvector()):
                 # Calculate the reciprocal lattice vectors that intersect the Ewald sphere.
                 intersection, excitation_error = get_intersection_with_ewalds_sphere(
                     recip,
-                    rot,
+                    optical_axis,
                     wavelength,
                     max_excitation_error,
                     self.precession_angle,
@@ -407,7 +409,7 @@ def get_intersecting_reflections(
 
         intersection, excitation_error = get_intersection_with_ewalds_sphere(
             recip,
-            rot,
+            rot * Vector3d.zvector(),
             wavelength,
             max_excitation_error,
             self.precession_angle,
@@ -422,11 +424,10 @@ def get_intersecting_reflections(
         )
         return intersected_vectors, hkl, shape_factor
 
-
 # TODO consider refactoring into a seperate file
 def get_intersection_with_ewalds_sphere(
     recip: DiffractingVector,
-    rot: Rotation,
+    optical_axis: Vector3d,
     wavelength: float,
     max_excitation_error: float,
     precession_angle: float = 0,
@@ -437,8 +438,8 @@ def get_intersection_with_ewalds_sphere(
     ----------
     recip
         The reciprocal lattice vectors to rotate.
-    rot
-        The rotation to apply to the reciprocal lattice vectors.
+    optical_axis
+        Normalised vector representing the direction of the beam
     wavelength
         The wavelength of the electrons in Angstroms.
     max_excitation_error
@@ -456,15 +457,38 @@ def get_intersection_with_ewalds_sphere(
     excitation_error
         Excitation error of all vectors
     """
-    # Identify the excitation errors of all points (distance from point to Ewald sphere)
+    if precession_angle == 0:
+        return _get_intersection_with_ewalds_sphere_without_precession(
+            recip.data,
+            optical_axis.data.squeeze(),
+            wavelength,
+            max_excitation_error
+        )
+    return _get_intersection_with_ewalds_sphere_with_precession(
+        recip.data,
+        optical_axis.data.squeeze(),
+        wavelength,
+        max_excitation_error,
+        precession_angle
+    )
 
+@njit(
+    "float64[:](float64[:, :], float64[:], float64)",
+    fastmath=True,
+)
+def _calculate_excitation_error(
+    recip: np.ndarray,
+    optical_axis_vector: np.ndarray,
+    wavelength: float,
+) -> np.ndarray:
     # Instead of rotating vectors, rotate Ewald's sphere to find intersections.
     # Only rotate the intersecting vectors.
     # Using notation from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
     r = 1 / wavelength
-    u = rot.to_matrix().squeeze() @ np.array([0, 0, 1])
+    # u = rot @ np.array([0, 0, 1])
+    u = optical_axis_vector
     c = r * u
-    o = recip.data
+    o = recip
 
     diff = o - c
     dot = np.dot(u, diff.T)
@@ -475,53 +499,85 @@ def get_intersection_with_ewalds_sphere(
     sqrt_nabla = np.sqrt(nabla)
     d = -dot - sqrt_nabla
 
-    excitation_error = d
+    return d
 
-    # determine the pre-selection reflections
-    if precession_angle == 0:
-        intersection = np.abs(excitation_error) < max_excitation_error
-    else:
-        # Using the following script to find equations for upper and lower bounds for precessing Ewald's sphere
-        """
-        import sympy
-        import numpy as np
-
-        a = sympy.Symbol("a") # Precession angle
-        r = sympy.Symbol("r") # Ewald's sphere radius
-        rho, z = sympy.symbols("rho z") # cylindrical coordinates of reflection
-
-        rot = lambda ang: np.asarray([[sympy.cos(ang), -sympy.sin(ang)],[sympy.sin(ang), sympy.cos(ang)]])
-
-        u = np.asarray([0, 1])
-        c = r * u
-        cl = rot(a) @ c
-        cr = rot(-a) @ c
-        o = np.asarray([rho, z])
-
-        def get_d(_c):
-            diff = o - _c
-            dot = np.dot(u, diff)
-            nabla = dot**2 - sum(i**2 for i in diff) + r**2
-            sqrt_nabla = nabla**0.5
-            return  -dot - sqrt_nabla
-
-        d = get_d(c)
-        d_upper = get_d(cl)
-        d_lower = get_d(cr)
-
-        print(d.simplify())             # r - z - (r**2 - rho**2)**0.5
-        print((d_upper - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) + rho)**2)**0.5
-        print((d_lower - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) - rho)**2)**0.5
-        """
-        # In the above script, d is the same as before.
-        # We need the distance of the reflections from the incident beam, i.e. the cylindrical coordinate rho
-        # (using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation):
-        rho = np.linalg.norm(np.dot(o, u)[:, np.newaxis] * u - o, axis=1)
-        a = np.deg2rad(precession_angle)
-        first_half = r * np.cos(a) - r + (r**2 - rho**2) ** 0.5
-        upper = first_half - (r**2 - (r * np.sin(a) + rho) ** 2) ** 0.5
-        lower = first_half - (r**2 - (r * np.sin(a) - rho) ** 2) ** 0.5
-        intersection = (d < (upper + max_excitation_error)) & (
-            d > (lower - max_excitation_error)
-        )
+@njit(
+    "Tuple((bool[:], float64[:]))(float64[:, :], float64[:], float64, float64)",
+    fastmath=True,
+)
+def _get_intersection_with_ewalds_sphere_without_precession(
+    recip: np.ndarray,
+    optical_axis_vector: np.ndarray,
+    wavelength: float,
+    max_excitation_error: float,
+) -> Tuple[np.ndarray, np.ndarray]:
+    excitation_error = _calculate_excitation_error(recip, optical_axis_vector, wavelength)
+    intersection = np.abs(excitation_error) < max_excitation_error
     return intersection, excitation_error
+
+    
+@njit(
+    "Tuple((bool[:], float64[:]))(float64[:, :], float64[:], float64, float64, float64)",
+    fastmath=True,
+)
+def _get_intersection_with_ewalds_sphere_with_precession(
+    recip: np.ndarray,
+    optical_axis_vector: np.ndarray,
+    wavelength: float,
+    max_excitation_error: float,
+    precession_angle: float,
+) -> Tuple[np.ndarray, np.ndarray]:
+    # Using the following script to find equations for upper and lower bounds for precessing Ewald's sphere
+    # (names are same as in _get_excitation_error_no_precession)
+    """
+    import sympy
+    import numpy as np
+
+    a = sympy.Symbol("a") # Precession angle
+    r = sympy.Symbol("r") # Ewald's sphere radius
+    rho, z = sympy.symbols("rho z") # cylindrical coordinates of reflection
+
+    rot = lambda ang: np.asarray([[sympy.cos(ang), -sympy.sin(ang)],[sympy.sin(ang), sympy.cos(ang)]])
+
+    u = np.asarray([0, 1])
+    c = r * u
+    cl = rot(a) @ c
+    cr = rot(-a) @ c
+    o = np.asarray([rho, z])
+
+    def get_d(_c):
+        diff = o - _c
+        dot = np.dot(u, diff)
+        nabla = dot**2 - sum(i**2 for i in diff) + r**2
+        sqrt_nabla = nabla**0.5
+        return  -dot - sqrt_nabla
+
+    d = get_d(c)
+    d_upper = get_d(cl)
+    d_lower = get_d(cr)
+
+    print(d.simplify())             # r - z - (r**2 - rho**2)**0.5
+    print((d_upper - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) + rho)**2)**0.5
+    print((d_lower - d).simplify()) # r*cos(a) - r + (r**2 - rho**2)**0.5 - (r**2 - (r*sin(a) - rho)**2)**0.5
+    """
+    d = _calculate_excitation_error(recip, optical_axis_vector, wavelength)
+
+    r = 1 / wavelength
+    u = optical_axis_vector
+    o = recip
+
+    excitation_error = d
+    # We need the distance of the reflections from the incident beam, i.e. the cylindrical coordinate rho
+    # (using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation):
+
+    # Numba does not support norm with axes, implement manually
+    rho = np.sum((np.dot(o, u)[:, np.newaxis] * u - o)**2, axis=1)**0.5
+    a = np.deg2rad(precession_angle)
+    first_half = r * np.cos(a) - r + (r**2 - rho**2) ** 0.5
+    upper = first_half - (r**2 - (r * np.sin(a) + rho) ** 2) ** 0.5
+    lower = first_half - (r**2 - (r * np.sin(a) - rho) ** 2) ** 0.5
+    intersection = (d < (upper + max_excitation_error)) & (
+        d > (lower - max_excitation_error)
+    )
+    return intersection, excitation_error
+

From 8d2c0f3888962609fd46b6ce4aa55500726b7715 Mon Sep 17 00:00:00 2001
From: Viljar Femoen <viljar.femoen@hotmail.no>
Date: Mon, 27 Jan 2025 16:24:19 +0100
Subject: [PATCH 17/17] Formatting

---
 .../crystallography/_diffracting_vector.py    |  1 -
 diffsims/generators/simulation_generator.py   | 19 ++++++++++---------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/diffsims/crystallography/_diffracting_vector.py b/diffsims/crystallography/_diffracting_vector.py
index 85961645..2ab2f54f 100644
--- a/diffsims/crystallography/_diffracting_vector.py
+++ b/diffsims/crystallography/_diffracting_vector.py
@@ -20,7 +20,6 @@
 import numpy as np
 from diffpy.structure import Structure
 from orix.crystal_map import Phase
-from orix.vector.miller import _transform_space
 from orix.quaternion import Rotation
 
 
diff --git a/diffsims/generators/simulation_generator.py b/diffsims/generators/simulation_generator.py
index ad5cb8df..d454690d 100644
--- a/diffsims/generators/simulation_generator.py
+++ b/diffsims/generators/simulation_generator.py
@@ -424,6 +424,7 @@ def get_intersecting_reflections(
         )
         return intersected_vectors, hkl, shape_factor
 
+
 # TODO consider refactoring into a seperate file
 def get_intersection_with_ewalds_sphere(
     recip: DiffractingVector,
@@ -459,19 +460,17 @@ def get_intersection_with_ewalds_sphere(
     """
     if precession_angle == 0:
         return _get_intersection_with_ewalds_sphere_without_precession(
-            recip.data,
-            optical_axis.data.squeeze(),
-            wavelength,
-            max_excitation_error
+            recip.data, optical_axis.data.squeeze(), wavelength, max_excitation_error
         )
     return _get_intersection_with_ewalds_sphere_with_precession(
         recip.data,
         optical_axis.data.squeeze(),
         wavelength,
         max_excitation_error,
-        precession_angle
+        precession_angle,
     )
 
+
 @njit(
     "float64[:](float64[:, :], float64[:], float64)",
     fastmath=True,
@@ -501,6 +500,7 @@ def _calculate_excitation_error(
 
     return d
 
+
 @njit(
     "Tuple((bool[:], float64[:]))(float64[:, :], float64[:], float64, float64)",
     fastmath=True,
@@ -511,11 +511,13 @@ def _get_intersection_with_ewalds_sphere_without_precession(
     wavelength: float,
     max_excitation_error: float,
 ) -> Tuple[np.ndarray, np.ndarray]:
-    excitation_error = _calculate_excitation_error(recip, optical_axis_vector, wavelength)
+    excitation_error = _calculate_excitation_error(
+        recip, optical_axis_vector, wavelength
+    )
     intersection = np.abs(excitation_error) < max_excitation_error
     return intersection, excitation_error
 
-    
+
 @njit(
     "Tuple((bool[:], float64[:]))(float64[:, :], float64[:], float64, float64, float64)",
     fastmath=True,
@@ -571,7 +573,7 @@ def get_d(_c):
     # (using https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation):
 
     # Numba does not support norm with axes, implement manually
-    rho = np.sum((np.dot(o, u)[:, np.newaxis] * u - o)**2, axis=1)**0.5
+    rho = np.sum((np.dot(o, u)[:, np.newaxis] * u - o) ** 2, axis=1) ** 0.5
     a = np.deg2rad(precession_angle)
     first_half = r * np.cos(a) - r + (r**2 - rho**2) ** 0.5
     upper = first_half - (r**2 - (r * np.sin(a) + rho) ** 2) ** 0.5
@@ -580,4 +582,3 @@ def get_d(_c):
         d > (lower - max_excitation_error)
     )
     return intersection, excitation_error
-