Speed up simulations

diffsims/crystallography/_diffracting_vector.py

@@ -18,10 +18,40 @@
 
 from diffsims.crystallography import ReciprocalLatticeVector
 import numpy as np
-from orix.vector.miller import _transform_space
+from diffpy.structure import Structure
+from orix.crystal_map import Phase
 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):
+        # 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
+        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.
@@ -93,27 +123,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
 
@@ -151,14 +163,11 @@ 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)
+        return self.__class__(new_phase, xyz=vecs.data)
 
     @property
     def intensity(self):
@@ -177,11 +186,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"""

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
 import numba as nb
 import numpy as np
 from orix.vector import Miller, Vector3d
@@ -644,7 +642,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 +653,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
         --------
@@ -692,32 +694,17 @@ def calculate_structure_factor(self, scattering_params="xtables"):
         """
 
         # 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,
             g_hkls_array=self.phase.structure.lattice.rnorm(hkl_unique),
+            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`
@@ -930,16 +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()
 
-        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.
 
@@ -1506,7 +1489,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
@@ -1517,48 +1500,59 @@ 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
         -------
         ReciprocalLatticeVector
             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, ignore_zero=False
+            )
 
             if use_symmetry:
                 operations = miller.phase.point_group
                 n_v = v.size
                 v2 = operations.outer(v).flatten().reshape(*(n_v, operations.size))
-                data = v2.data.round(6)  # Reduced precision
+                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]
-                _, idx = np.unique(data_sorted, return_index=True, axis=0)
-                v = v[idx[::-1]]
+                    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)
             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
 
@@ -1596,62 +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, atom2.xyz, 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
-
-    """
-
-    new_structure = Structure(lattice=structure.lattice)
-
-    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)
-
-    return new_structure
-
-
 @nb.njit(
     "int64[:](float64[:, :], float64[:, :], int64[:])",
     cache=True,