Add elastic line fitting to centroiding algorithm

src/toolbox_scs/detectors/hrixs.py

@@ -302,6 +302,13 @@ class hRIXS:
     DARK_MASK_THRESHOLD = 100
     MASK_AVG_X = np.s_[1850:2000]
     MASK_AVG_Y = np.s_[500:1500]
+    
+    # Early energy calibration from runs 31-36 (to estimate mono drift)
+    PIX2EV_POLY = [6.31196512e-02, 6.05502748e+02]
+    # Width of fit window to calculate hv (estimated mono ± this)
+    HV_FIT_DELTA = 30
+    # Gaussian fit parameters for elastic line finding
+    HV_FIT_SIGMA = 2
 
     ENERGY_INTERCEPT = 0
     ENERGY_SLOPE = 1
@@ -319,6 +326,23 @@ class hRIXS:
                       'factor', 'range', 'bins',
                       'method', 'fields')
         return {param: getattr(self, param.upper()) for param in params}
+    
+    def pixel2energy(self, pixel):         
+        # Calculates energy based on pixel position from PIX2EV_POLY (re-fit as needed).
+        return np.polyval(self.PIX2EV_POLY, pixel)
+
+    def energy2pixel(self, ev):
+        # Calculates the expected pixel position of a given energy.
+        # Only works (somewhat) predictably for 1st and 2nd order polynomials!
+        if len(self.PIX2EV_POLY)>3:
+            # The length of PIX2EV_POLY is order+1
+            raise ValueError('Too high polynomial order!')
+        epoly = np.poly1d(self.PIX2EV_POLY)
+        pixel = (epoly-ev).roots[-1]
+        if np.imag(pixel) == 0:
+            return pixel
+        else:
+            raise ValueError('Complex root! (Out of fit bounds?)')
 
     def from_run(self, runNB, proposal=None, extra_fields=()):
         if proposal is None:
@@ -393,7 +417,8 @@ class hRIXS:
         return self.CURVE_A, self.CURVE_B
 
     
-    def centroid(self, data, bins=None, return_hits=False):
+    def centroid(self, data, bins=None, return_hits=False, fit_elastic=False,
+                 hv=None, fit_delta=None):
         #*************************************************************
         # Carry out hit finding on data and bin them in grid
         # Allows for 
@@ -416,10 +441,30 @@ class hRIXS:
         hit_y = []
         hits = []
         ret = np.zeros((len(data["hRIXS_det"]), bins))
+        #*************************************************************
+        # If the 'actual' photon energy is desired, we'll need to know
+        # where to start looking for the elastic line
+        #*************************************************************
+        if fit_elastic:
+            elastic_pixel = np.zeros(len(data["hRIXS_det"]))
+            found_hv = np.zeros(len(data["hRIXS_det"]))
+            if fit_delta == None:
+                fit_delta = self.HV_FIT_DELTA
+            if hv == None:
+                # ?? We'll make a guess later
+                hv = -1 * np.ones(len(data["hRIXS_det"]))
+            else:
+                # hv is given as one value or a train-vector
+                if type(hv) == type([]):
+                    if len(hv) != len(data["hRIXS_det"]):
+                        raise ValueError("Size of hv vector doesn't match trainID")
+                else:
+                    hv = hv * np.ones(len(data["hRIXS_det"]))
+                
         #*************************************************************
         # Handle each Aq image separately
         #*************************************************************
-        for image, r in zip(data["hRIXS_det"], ret):
+        for i, (image, r) in enumerate(zip(data["hRIXS_det"], ret)):
             use_image = image.to_numpy()
             #*************************************************************
             # Treat background by optionally 
@@ -473,6 +518,25 @@ class hRIXS:
                 rc[:, 0], bins=bins,
                 range=(0, self.Y_RANGE.stop - self.Y_RANGE.start))
             r[:] = hy
+            #*************************************************************
+            # Return the found elastic peak energy, if requested
+            #*************************************************************
+            if fit_elastic:
+                if hv[i] == -1:
+                    # Assume elastic line to be the strongest peak
+                    gx = np.where(hy == max(hy))[0][0]
+                else:
+                    # Use given guess for where it is
+                    gx = np.where(hx + self.Y_RANGE.start > self.energy2pixel(hv[i]))[0][0]
+                popt, pcov = curve_fit(gauss1d, hx[gx-fit_delta:gx+fit_delta],
+                                       hy[gx-fit_delta:gx+fit_delta],
+                                      p0 = [max(hx[gx-fit_delta:gx+fit_delta]),
+                                           hx[gx], self.HV_FIT_SIGMA, 0],
+                                      bounds=([0, 0, 0, -np.Inf],[np.Inf, np.Inf, 100, np.Inf]))
+                
+                elastic_pixel[i] = popt[1]+ self.Y_RANGE.start
+                found_hv[i] = self.pixel2energy(popt[1]+ self.Y_RANGE.start)
+                                       
         #*************************************************************
         # Setup and assing a linear energy grid
         #*************************************************************
@@ -487,6 +551,12 @@ class hRIXS:
                         xhits=(("trainId"), hit_x),
                         yhits=(("trainId"), hit_y))
         #**********************************************
+        # If hv was fitted, return it
+        #**********************************************
+        if fit_elastic:
+            data = data.assign(fitted_hv = (("trainId"), found_hv),
+                              elastic_pixel = (("trainId"), elastic_pixel))
+        #**********************************************
         # Always assign the spectrum to data
         #**********************************************
         data = data.assign(spectrum=(("trainId", "energy"), ret))