Jit some structure factor functions

Author: viljarjf

diffsims/utils/sim_utils.py

@@ -19,7 +19,7 @@
 import collections
 import math
 
-import diffpy.structure
+from numba import njit
 import numpy as np
 from scipy.constants import h, m_e, e, c, pi, mu_0
 
@@ -215,7 +215,7 @@ def get_vectorized_list_for_atomic_scattering_factors(
 
     return coeffs, fcoords, occus, dwfactors
 
-
+@njit(fastmath=True)
 def get_atomic_scattering_factors(g_hkl_sq, coeffs, scattering_params):
     """Calculate atomic scattering factors for n atoms.
 
@@ -277,9 +277,32 @@ def _get_kinematical_structure_factor(
         # Set all atomic scattering factors to 1
         atomic_scattering_factor = np.ones((gspacing_squared.shape[0], coeffs.shape[0]))
 
+    return _numba_get_kinematical_structure_factor(
+        structure.lattice.stdbase,
+        structure.lattice.recbase,
+        xyz,
+        atomic_scattering_factor,
+        occupancy,
+        g_indices,
+        gspacing_squared,
+        dwfactors,
+    )
+
+
+@njit(fastmath=True)
+def _numba_get_kinematical_structure_factor(
+        stdbase: np.ndarray,
+        recbase: np.ndarray,
+        xyz: np.ndarray,
+        atomic_scattering_factor: np.ndarray,
+        occupancy: np.ndarray,
+        g_indices: np.ndarray,
+        g_spacing_squared: np.ndarray,
+        dwfactors: np.ndarray,
+) -> np.ndarray:
     # Express the atom positions in the same reference frame as the
     # Miller indices
-    mat = np.linalg.inv(np.dot(structure.lattice.stdbase, structure.lattice.recbase))
+    mat = np.linalg.inv(np.dot(stdbase, recbase))
     xyz = np.dot(xyz, mat)
 
     # Calculate the complex structure factor
@@ -288,11 +311,10 @@ def _get_kinematical_structure_factor(
         * occupancy
         * np.exp(
             2j * np.pi * np.dot(g_indices, xyz.T)
-            - 0.25 * np.outer(gspacing_squared, dwfactors)
+            - 0.25 * np.outer(g_spacing_squared, dwfactors)
         ),
         axis=-1,
     )
-
     return structure_factor