adds prominence

Former-commit-id: 8054c2f
bsavitzky · Feb 7, 2024 · 36b4fe2 · 36b4fe2
1 parent 7f053fe
commit 36b4fe2
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 40 deletions.
diff --git a/py4DSTEM/datacube/diskdetection/braggfinder.py b/py4DSTEM/datacube/diskdetection/braggfinder.py
@@ -338,9 +338,10 @@ def find_bragg_vectors(
             'edge' : 0,
             'sigma' : 0,
             'n_peaks_max' : 10000,
+            'min_prominence' : 0,
+            'prominence_kernel_size' : 3,
             'min_rel_intensity' : 0,
             'ref_peak' : 0,
-            'method' : "prominence",
         }
         corr = corr_defaults if corr is None else corr_defaults | corr
         thresh = thresh_defaults if thresh is None else thresh_defaults | thresh
@@ -409,46 +410,40 @@ def _cross_correlate(dp):
             return cc
 
         # threshold
-        method_options = [
-            'prominence',
-            'absolute'
-        ]
-        assert(thresh['method'] in method_options)
         def _threshold(cc):
             """ vec = _threshold(cc)
             """
-            if thresh['method'] == 'prominence':
-                # TODO
-                raise Exception("prominence method has not been implemented yet")
-
+            cc_real = np.maximum(np.real(np.fft.ifft2(cc)),0)
+            if thresh['subpixel'] == 'multicorr':
+                return get_maxima_2D(
+                    cc_real,
+                    subpixel='multicorr',
+                    upsample_factor=thresh['upsample_factor'],
+                    sigma=thresh['sigma'],
+                    minAbsoluteIntensity=thresh['min_intensity'],
+                    minProminence=thresh['min_prominence'],
+                    prominenceKernelSize=thresh['prominence_kernel_size'],
+                    minRelativeIntensity=thresh['min_rel_intensity'],
+                    relativeToPeak=thresh['ref_peak'],
+                    minSpacing=thresh['min_spacing'],
+                    edgeBoundary=thresh['edge'],
+                    maxNumPeaks=thresh['n_peaks_max'],
+                    _ar_FT=cc,
+                )
             else:
-                cc_real = np.maximum(np.real(np.fft.ifft2(cc)),0)
-                if thresh['subpixel'] == 'multicorr':
-                    return get_maxima_2D(
-                        cc_real,
-                        subpixel='multicorr',
-                        upsample_factor=thresh['upsample_factor'],
-                        sigma=thresh['sigma'],
-                        minAbsoluteIntensity=thresh['min_intensity'],
-                        minRelativeIntensity=thresh['min_rel_intensity'],
-                        relativeToPeak=thresh['ref_peak'],
-                        minSpacing=thresh['min_spacing'],
-                        edgeBoundary=thresh['edge'],
-                        maxNumPeaks=thresh['n_peaks_max'],
-                        _ar_FT=cc,
-                    )
-                else:
-                    return get_maxima_2D(
-                        cc_real,
-                        subpixel=thresh['subpixel'],
-                        sigma=thresh['sigma'],
-                        minAbsoluteIntensity=thresh['min_intensity'],
-                        minRelativeIntensity=thresh['min_rel_intensity'],
-                        relativeToPeak=thresh['ref_peak'],
-                        minSpacing=thresh['min_spacing'],
-                        edgeBoundary=thresh['edge'],
-                        maxNumPeaks=thresh['n_peaks_max'],
-                    )
+                return get_maxima_2D(
+                    cc_real,
+                    subpixel=thresh['subpixel'],
+                    sigma=thresh['sigma'],
+                    minAbsoluteIntensity=thresh['min_intensity'],
+                    minProminence=thresh['min_prominence'],
+                    prominenceKernelSize=thresh['prominence_kernel_size'],
+                    minRelativeIntensity=thresh['min_rel_intensity'],
+                    relativeToPeak=thresh['ref_peak'],
+                    minSpacing=thresh['min_spacing'],
+                    edgeBoundary=thresh['edge'],
+                    maxNumPeaks=thresh['n_peaks_max'],
+                )
 
 
         # prepare the template

diff --git a/py4DSTEM/utils/get_maxima.py b/py4DSTEM/utils/get_maxima.py
@@ -1,5 +1,6 @@
 import numpy as np
 from scipy.ndimage import gaussian_filter
+from scipy.signal import medfilt
 try:
     import cupy as cp
 except (ModuleNotFoundError, ImportError):
@@ -75,18 +76,26 @@ def get_maxima_1D(ar, sigma=0, minSpacing=0, minRelativeIntensity=0, relativeToP
 def filter_2D_maxima(
     maxima,
     minAbsoluteIntensity=0,
+    minProminence=0,
+    prominenceKernelSize=3,
     minRelativeIntensity=0,
     relativeToPeak=0,
     minSpacing=0,
     edgeBoundary=1,
     maxNumPeaks=1,
+    _ar = None,
 ):
     """
     Parameters
     ----------
     maxima : a numpy structured array with fields 'x', 'y', 'intensity'
     minAbsoluteIntensity : number
         delete counts with intensity below this value
+    minProminence : number
+        delete counts whose intensity above the local median is
+        below this value. _ar must be passed with this arg
+    prominenceKernelSize : odd number
+        kernel size for prominence local median comparison
     minRelativeIntensity : number
         delete counts with intensity below this value times
         the intensity of the i'th peak, where i is given by `relativeToPeak`
@@ -99,6 +108,8 @@ def filter_2D_maxima(
         delete peaks within this distance of the image edge
     maxNumPeaks : int
         maximum number of peaks to return
+    _ar : array or None
+        if minProminence is passed, this must be the array
 
     Returns
     -------
@@ -110,6 +121,17 @@ def filter_2D_maxima(
         deletemask = maxima["intensity"] < minAbsoluteIntensity
         maxima = maxima[~deletemask]
 
+    # Remove maxima which are too dim relative to local median
+    if minProminence > 0:
+        assert(_ar is not None), "Array for median filter wasn't passed"
+        med = medfilt(_ar,prominenceKernelSize)
+        compare = maxima["intensity"] - med[
+            maxima['x'].astype(int),
+            maxima['y'].astype(int)
+        ]
+        deletemask = compare < minProminence
+        maxima = maxima[~deletemask]
+
     # Remove maxima which are too dim, compared to the n-th brightest
     if (minRelativeIntensity > 0) & (len(maxima) > relativeToPeak):
         assert isinstance(relativeToPeak, (int, np.integer))
@@ -141,15 +163,14 @@ def filter_2D_maxima(
 
 
 
-
-
-
 def get_maxima_2D(
     ar,
     subpixel="poly",
     upsample_factor=16,
     sigma=0,
     minAbsoluteIntensity=0,
+    minProminence=0,
+    prominenceKernelSize=3,
     minRelativeIntensity=0,
     relativeToPeak=0,
     minSpacing=0,
@@ -177,6 +198,11 @@ def get_maxima_2D(
         the maximum number of maxima to return
     minAbsoluteIntensity : number
         minimum intensity threshold
+    minProminence : number
+        intensity threshold, in absolute units, above the local
+        median
+    prominenceKernelSize : odd number
+        kernel size for prominence local median comparison
     minRelativeIntensity : number
         intensity threshold relative to the N'th brightest peak,
         where N is given by `relativeToPeak`
@@ -241,11 +267,14 @@ def get_maxima_2D(
     maxima = filter_2D_maxima(
         maxima,
         minAbsoluteIntensity=minAbsoluteIntensity,
+        minProminence=minProminence,
+        prominenceKernelSize=prominenceKernelSize,
         minRelativeIntensity=minRelativeIntensity,
         relativeToPeak=relativeToPeak,
         minSpacing=minSpacing,
         edgeBoundary=edgeBoundary,
         maxNumPeaks=maxNumPeaks,
+        _ar = ar
     )
 
     if subpixel == "pixel":