Merge pull request #552 from sezelt/robust-fix

Add robust fitting back to `fit_origin`
py4dstem · Nov 1, 2023 · 02c419a · 02c419a
2 parents af280c3 + fac36a7
commit 02c419a
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/py4DSTEM/braggvectors/braggvector_methods.py b/py4DSTEM/braggvectors/braggvector_methods.py
@@ -552,14 +552,20 @@ def fit_origin(
         from py4DSTEM.process.calibration import fit_origin
 
         if mask_check_data is True:
-            # TODO - replace this bad hack for the mask for the origin fit
-            mask = np.logical_not(q_meas[0] == 0)
-            qx0_fit, qy0_fit, qx0_residuals, qy0_residuals = fit_origin(
-                tuple(q_meas),
-                mask=mask,
-            )
-        else:
-            qx0_fit, qy0_fit, qx0_residuals, qy0_residuals = fit_origin(tuple(q_meas))
+            data_mask = np.logical_not(q_meas[0] == 0)
+            if mask is None:
+                mask = data_mask
+            else:
+                mask = np.logical_and(mask, data_mask)
+
+        qx0_fit, qy0_fit, qx0_residuals, qy0_residuals = fit_origin(
+            tuple(q_meas),
+            mask=mask,
+            fitfunction=fitfunction,
+            robust=robust,
+            robust_steps=robust_steps,
+            robust_thresh=robust_thresh,
+        )
 
         # try to add to calibration
         try: