From da0e897156536651282c9959f793b428aa63da76 Mon Sep 17 00:00:00 2001
From: bsavitzky <ben.savitzky@gmail.com>
Date: Thu, 9 May 2024 12:00:04 -0400
Subject: [PATCH] fix probe norm bug

---
 py4DSTEM/braggvectors/probe.py | 95 +++++++++++++++++++++++++++++++---
 1 file changed, 89 insertions(+), 6 deletions(-)

diff --git a/py4DSTEM/braggvectors/probe.py b/py4DSTEM/braggvectors/probe.py
index 464c2f2a4..f7340b69b 100644
--- a/py4DSTEM/braggvectors/probe.py
+++ b/py4DSTEM/braggvectors/probe.py
@@ -4,7 +4,9 @@
 from typing import Optional
 from warnings import warn
 
+from emdfile import Metadata
 from py4DSTEM.data import DiffractionSlice, Data
+from py4DSTEM.visualize import show
 from scipy.ndimage import binary_opening, binary_dilation, distance_transform_edt
 
 
@@ -43,12 +45,16 @@ def __init__(self, data: np.ndarray, name: Optional[str] = "probe"):
             self, name=name, data=data, slicelabels=["probe", "kernel"]
         )
 
+        # initialize metadata params
+        self.metadata = Metadata(name='params')
+        self.alpha = None
+        self.origin = None
+
     ## properties
 
     @property
     def probe(self):
         return self.get_slice("probe").data
-
     @probe.setter
     def probe(self, x):
         assert x.shape == (self.data.shape[1:])
@@ -57,12 +63,25 @@ def probe(self, x):
     @property
     def kernel(self):
         return self.get_slice("kernel").data
-
     @kernel.setter
     def kernel(self, x):
         assert x.shape == (self.data.shape[1:])
         self.data[1, :, :] = x
 
+    @property
+    def alpha(self):
+        return self.metadata['params']['alpha']
+    @alpha.setter
+    def alpha(self, x):
+        self.metadata['params']['alpha'] = x
+
+    @property
+    def origin(self):
+        return self.metadata['params']['origin']
+    @origin.setter
+    def origin(self, x):
+        self.metadata['params']['origin'] = x
+
     # read
     @classmethod
     def _get_constructor_args(cls, group):
@@ -181,8 +200,11 @@ def measure_disk(
         thresh_lower=0.01,
         thresh_upper=0.99,
         N=100,
-        returncalc=True,
         data=None,
+        zero_vacuum=True,
+        alpha_max=1.2,
+        returncalc=True,
+        plot=True,
     ):
         """
         Finds the center and radius of an average probe image.
@@ -207,12 +229,20 @@ def measure_disk(
             the upper limit of threshold values
         N : int
             the number of thresholds / masks to use
-        returncalc : True
-            toggles returning the answer
         data : 2d array, optional
             if passed, uses this 2D array in place of the probe image when
             performing the computation. This also supresses storing the
             results in the Probe's calibration metadata
+        zero_vacuum : bool
+            if True, sets pixels beyond alpha_max * the semiconvergence angle
+            to zero. Ignored if `data` is not None
+        alpha_max : number
+            sets the maximum scattering angle in the probe image, beyond which
+            values are set to zero if `zero_vacuum` is True
+        returncalc : True
+            toggles returning the answer
+        plot : bool
+            toggles visualizing results
 
         Returns
         -------
@@ -244,9 +274,11 @@ def measure_disk(
         mask = im > immax * thresh
         x0, y0 = get_CoM(im * mask)
 
-        # Store metadata and return
+        # Store metadata
         ans = r, x0, y0
         if data is None:
+            self.alpha = r
+            self.origin = (x0,y0)
             try:
                 self.calibration.set_probe_param(ans)
             except AttributeError:
@@ -254,9 +286,60 @@ def measure_disk(
                     f"Couldn't store the probe parameters in metadata as no calibration was found for this Probe instance, {self}"
                 )
                 pass
+
+        if data is None and zero_vacuum:
+            self.zero_vacuum( alpha_max=alpha_max)
+
+        # show result
+        if plot:
+            show(
+                im,
+                circle={
+                    'center' : (x0,y0),
+                    'R' : r,
+                    'fill' : True,
+                    'alpha' : 0.36
+                }
+            )
+
+        # return
         if returncalc:
             return ans
 
+
+    def zero_vacuum(
+        self,
+        alpha_max = 1.2,
+    ):
+        """
+        Sets pixels outside of the probe's central disk to zero.
+
+        The probe origin and convergence semiangle must be set for this
+        method to run - these can be set using `measure_disk`. Pixels are
+        defined as outside the central disk if their distance from the origin
+        exceeds the semiconvergence angle * alpha_max.
+
+        Parameters
+        ----------
+        alpha_max : number
+            Pixels farther than this number times the semiconvergence angle
+            from the origin are set to zero
+        """
+        # validate inputs
+        assert(self.alpha is not None), "no probe semiconvergence angle found; try running `Probe.measure_disk`"
+        assert(self.origin is not None), "no probe origin found; try running `Probe.measure_disk`"
+        # make a mask
+        qyy,qxx = np.meshgrid(
+            np.arange(self.shape[1]),
+            np.arange(self.shape[0])
+        )
+        qrr = np.hypot(qxx-self.origin[0],qyy-self.origin[1])
+        mask = qrr < self.alpha*alpha_max
+        # zero the vacuum
+        self.probe *= mask
+        pass
+
+
     # Kernel generation methods
 
     def get_kernel(