Initial commit with alignment function draft

src/toolbox_scs/detectors/hrixs.py

@@ -491,6 +491,103 @@ class hRIXS:
         #**********************************************
         data = data.assign(spectrum=(("trainId", "energy"), ret))
         return data
+
+    
+    def align_readouts(self,data,method,start,stop):
+        import scipy as sp
+        from scipy import optimize 
+        #********************************************
+        # aligns spectra in a given data xarray
+        # METHOD
+        # -max_value
+        # -autocorrelation
+        # -gauss_fit
+        # start and stop are values of data.energy
+        # that define the range of these operations
+        # RETURNS MAXIMUM POSITIONS AND data with 
+        # shifted spectra
+        #********************************************
+        searchinds = (data.energy >=start)*(data.energy <=stop)
+        peak_posis = []
+        #********************************************
+        # Simple maximum alignment
+        #********************************************
+        if method.lower() == 'max_value':
+            #********************************************
+            # Find the max for each of the spectra
+            #********************************************
+            for spec in data.spectrum:
+                x          = data.energy.to_numpy()[searchinds]
+                y          = spec.to_numpy()[searchinds]
+                maxipos    = np.argmax(y)
+                peak_posis.append(x[maxipos])
+        #********************************************
+        # Alignment based on autocorrelation
+        # this is a relative alignment method
+        #********************************************
+        elif method.lower() == 'autocorrelation':
+            #********************************************
+            # Find the max for each of the spectra
+            #********************************************
+            for ind,spec in enumerate(data.spectrum):
+                if ind == 0:
+                    x0       = data.energy.to_numpy()[searchinds]
+                    y0       = spec.to_numpy()[searchinds]
+                    maxipos0 = np.argmax(spec.to_numpy()[searchinds])
+                    peak_posis.append(x0[maxipos0])
+                else:        
+                    x        = data.energy.to_numpy()[searchinds]
+                    y        = spec.to_numpy()[searchinds]
+                    corr_len = np.sum(searchinds)
+                    corr     = sp.signal.correlate(y,y0, mode='full')
+                    maxpos   = np.argmax(corr)
+                    shift    = maxpos-corr_len
+                    peak_posis.append(x[maxipos0+shift])
+        elif method.lower() == 'gauss_fit':
+            #********************************************
+            # Define needed functions
+            #********************************************
+            def Gauss(grid,x0,sigma):
+                # Returns a normalized bell curve
+                # with center at x0 and sigma
+                # on grid
+                return 1.0/(sigma*np.sqrt(2.0*np.pi))*np.exp(-0.5*(grid-x0)**2/sigma**2)
+
+            def Cost(p,grid,spec):
+                return np.sum(np.square(p[0]*Gauss(grid,p[1],p[2])-spec))
+            #********************************************
+            # Find the max for each of the spectra
+            #********************************************
+            for spec in data.spectrum:
+                x    = data.energy.to_numpy()[searchinds]
+                y    = spec.to_numpy()[searchinds]
+                #********************************************
+                # Initial Guess and bounds
+                #********************************************
+                area = np.sum(y)
+                mean = np.average(x,weights=y)
+                std  = np.sqrt(np.average((x-mean)**2,weights=y/area))
+                p0         = [area, mean, std]
+                #********************************************
+                # Bounds
+                #********************************************
+                bnds       = [[0,None],[start,stop],[0,2*(stop-start)]]
+                #********************************************
+                # Fit by minimizing least squares error
+                #********************************************
+                p          = optimize.minimize(Cost,p0,args=(x,y),bounds=bnds,method='L-BFGS-B',tol=1e-6,
+                                            options={'disp':0,'maxiter':1000000})
+                if p.success:
+                    peak_posis.append(p.x[1])
+                else:
+                    plt.figure()
+                    plt.plot(x,y,'.')
+                    plt.plot(x,p.x[0]*Gauss(x,p.x[1],p.x[2]))
+                    raise Exception('align_readouts(): can not fit a gaussian to the data.')
+        else: 
+            raise Exception('align_readouts() did recognize the method.')
+        return peak_posis    
+    
     
     def integrate(self, data):
         bins = self.Y_RANGE.stop - self.Y_RANGE.start