inline the centroid into the hRIXS class

src/toolbox_scs/detectors/hrixs.py

@@ -171,64 +171,6 @@ THRESHOLD = 510  # pixel counts above which a hit candidate is assumed
 CURVE_A = 2.19042931e-02  # curvature parameters as determined elsewhere
 CURVE_B = -3.02191568e-07
 
-
-def _esrf_centroid(image, threshold=THRESHOLD, curvature=(CURVE_A, CURVE_B)):
-    gs = 2
-    base = image.mean()
-    cp = np.argwhere(image[gs // 2: -gs // 2, gs // 2: -gs // 2] > threshold) + np.array([gs // 2, gs // 2])
-    if len(cp) > 100000:
-        raise RuntimeError('Threshold too low or acquisition time too long')
-
-    res = []
-    for cy, cx in cp:
-        spot = image[cy - gs // 2: cy + gs // 2 + 1, cx - gs // 2: cx + gs // 2 + 1] - base
-        spot[spot < 0] = 0
-        if (spot > image[cy, cx]).sum() == 0:
-            mx = np.average(np.arange(cx - gs // 2, cx + gs // 2 + 1), weights=spot.sum(axis=0))
-            my = np.average(np.arange(cy - gs // 2, cy + gs // 2 + 1), weights=spot.sum(axis=1))
-            my -= (curvature[0] + curvature[1] * mx) * mx
-            res.append((my, mx))
-    return res
-
-
-def _new_centroid(image, threshold=THRESHOLD, std_threshold=3.5, curvature=(CURVE_A, CURVE_B)):
-    """find the position of photons with sub-pixel precision
-
-    A photon is supposed to have hit the detector if the intensity within a
-    2-by-2 square exceeds a threshold. In this case the position of the photon
-    is calculated as the center-of-mass in a 4-by-4 square.
-
-    Return the list of x,y coordinate pairs, corrected by the curvature.
-    """
-    base = image.mean()
-    corners = image[1:, 1:] + image[:-1, 1:] + image[1:, :-1] + image[:-1, :-1]
-    if threshold is None:
-        threshold = corners.mean() + std_threshold * corners.std()
-    middle = corners[1:-1, 1:-1]
-    candidates = (
-            (middle > threshold)
-            * (middle >= corners[:-2, 1:-1]) * (middle > corners[2:, 1:-1])
-            * (middle >= corners[1:-1, :-2]) * (middle > corners[1:-1, 2:])
-            * (middle >= corners[:-2, :-2]) * (middle > corners[2:, :-2])
-            * (middle >= corners[:-2, 2:]) * (middle > corners[2:, 2:]))
-    cp = np.argwhere(candidates)
-    if len(cp) > 10000:
-        raise RuntimeError(
-            "too many peaks, threshold too low or acquisition time too high")
-
-    res = []
-    for cy, cx in cp:
-        spot = image[cy: cy + 4, cx: cx + 4] - base
-        mx = np.average(np.arange(cx, cx + 4), weights=spot.sum(axis=0))
-        my = np.average(np.arange(cy, cy + 4), weights=spot.sum(axis=1))
-        my -= (curvature[0] + curvature[1] * mx) * mx
-        res.append((my, mx))
-    return res
-
-
-centroid = _new_centroid
-
-
 def decentroid(res):
     res = np.array(res)
     ret = np.zeros(shape=(res.max(axis=0) + 1).astype(int))
@@ -410,6 +352,42 @@ class hRIXS:
         self.CURVE_B, self.CURVE_A, *_ = args
         return self.CURVE_A, self.CURVE_B
 
+    def centroid_one(self, image):
+        """find the position of photons with sub-pixel precision
+
+        A photon is supposed to have hit the detector if the intensity within a
+        2-by-2 square exceeds a threshold. In this case the position of the photon
+        is calculated as the center-of-mass in a 4-by-4 square.
+
+        Return the list of x, y coordinate pairs, corrected by the curvature.
+        """
+        base = image.mean()
+        corners = image[1:, 1:] + image[:-1, 1:] \
+                  + image[1:, :-1] + image[:-1, :-1]
+        if self.THRESHOLD is None:
+            threshold = corners.mean() + self.STD_THRESHOLD * corners.std()
+        else:
+            threshold = self.THRESHOLD
+        middle = corners[1:-1, 1:-1]
+        candidates = (
+            (middle > threshold)
+            * (middle >= corners[:-2, 1:-1]) * (middle > corners[2:, 1:-1])
+            * (middle >= corners[1:-1, :-2]) * (middle > corners[1:-1, 2:])
+            * (middle >= corners[:-2, :-2]) * (middle > corners[2:, :-2])
+            * (middle >= corners[:-2, 2:]) * (middle > corners[2:, 2:]))
+        cp = np.argwhere(candidates)
+        if len(cp) > 10000:
+            raise RuntimeError(
+                "too many peaks, threshold low or acquisition time too high")
+
+        res = []
+        for cy, cx in cp:
+            spot = image[cy: cy + 4, cx: cx + 4] - base
+            mx = np.average(np.arange(cx, cx + 4), weights=spot.sum(axis=0))
+            my = np.average(np.arange(cy, cy + 4), weights=spot.sum(axis=1))
+            res.append((mx, my))
+        return res
+
     def centroid(self, data, bins=None):
         """calculate a spectrum by finding the centroid of individual photons
 
@@ -427,23 +405,20 @@ class hRIXS:
             bins = self.BINS
         ret = np.zeros((len(data["hRIXS_det"]), bins))
         for image, r in zip(data["hRIXS_det"], ret):
-            c = centroid(
-                image.values[self.X_RANGE, self.Y_RANGE].T,
-                threshold=self.THRESHOLD,
-                std_threshold=self.STD_THRESHOLD,
-                curvature=(self.CURVE_A, self.CURVE_B))
+            c = self.centroid_one(
+                image.values[self.X_RANGE, self.Y_RANGE])
             if not len(c):
                 continue
             rc = np.array(c)
-            hy, hx = np.histogram(
-                rc[:, 0], bins=bins,
-                range=(0, self.Y_RANGE.stop - self.Y_RANGE.start))
-            r[:] = hy
+            r[:], _ = np.histogram(
+                rc[:, 0] - self.parabola(rc[:, 1]),
+                bins=bins, range=(0, self.Y_RANGE.stop - self.Y_RANGE.start))
 
-        data = data.assign_coords(
-            energy=np.linspace(self.Y_RANGE.start, self.Y_RANGE.stop, bins)
+        data.coords["energy"] = (
+            np.linspace(self.Y_RANGE.start, self.Y_RANGE.stop, bins)
             * self.ENERGY_SLOPE + self.ENERGY_INTERCEPT)
-        return data.assign(spectrum=(("trainId", "energy"), ret))
+        data['spectrum'] = (("trainId", "energy"), ret)
+        return data
 
     def parabola(self, x):
         return (self.CURVE_B * x + self.CURVE_A) * x