From 924f0e157fbe074bb972c69dc634d2084003619a Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Tue, 12 Dec 2023 11:06:42 -0500
Subject: [PATCH] add occupancy to structure factors

---
 py4DSTEM/process/diffraction/crystal.py       | 12 ++++++++----
 py4DSTEM/process/diffraction/crystal_bloch.py |  5 ++---
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index a2bf36a02..aa1eb8555 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -662,10 +662,14 @@ def calculate_structure_factors(
         # Calculate structure factors
         self.struct_factors = np.zeros(np.size(self.g_vec_leng, 0), dtype="complex64")
         for a0 in range(self.positions.shape[0]):
-            self.struct_factors += f_all[:, a0] * np.exp(
-                (2j * np.pi * self.occupancy[a0])
-                * np.sum(
-                    self.hkl * np.expand_dims(self.positions[a0, :], axis=1), axis=0
+            self.struct_factors += (
+                f_all[:, a0]
+                * self.occupancy[a0]
+                * np.exp(
+                    (2j * np.pi)
+                    * np.sum(
+                        self.hkl * np.expand_dims(self.positions[a0, :], axis=1), axis=0
+                    )
                 )
             )
 
diff --git a/py4DSTEM/process/diffraction/crystal_bloch.py b/py4DSTEM/process/diffraction/crystal_bloch.py
index 6a3c9b1ac..0878feb17 100644
--- a/py4DSTEM/process/diffraction/crystal_bloch.py
+++ b/py4DSTEM/process/diffraction/crystal_bloch.py
@@ -27,7 +27,6 @@ def calculate_dynamical_structure_factors(
     tol_structure_factor: float = 0.0,
     recompute_kinematic_structure_factors=True,
     g_vec_precision=None,
-    verbose=True,
 ):
     """
     Calculate and store the relativistic corrected structure factors used for Bloch computations
@@ -92,7 +91,7 @@ def calculate_dynamical_structure_factors(
 
     # Calculate the reciprocal lattice points to include based on k_max
 
-    k_max = np.asarray(k_max)
+    k_max: np.ndarray = np.asarray(k_max)
 
     if recompute_kinematic_structure_factors:
         if hasattr(self, "struct_factors"):
@@ -215,7 +214,7 @@ def get_f_e(q, Z, thermal_sigma, method):
 
     # Calculate structure factors
     struct_factors = np.sum(
-        f_e * np.exp(2.0j * np.pi * np.squeeze(self.positions[:, None, :] @ hkl)),
+        f_e * self.occupancy[:,None] * np.exp(2.0j * np.pi * np.squeeze(self.positions[:, None, :] @ hkl)),
         axis=0,
     )