find_offsets() added

src/cristallina/utils.py

@@ -7,6 +7,7 @@ import numpy as np
 from sfdata import SFDataFiles, sfdatafile, SFScanInfo, SFProcFile
 from xraydb import material_mu
 from joblib import Parallel, delayed, cpu_count
+from scipy.stats.stats import pearsonr
 
 def scan_info(run_number,base_path=None,small_data=True):
     """Returns SFScanInfo object for a given run number. 
@@ -35,7 +36,7 @@ def scan_from_run_number_or_scan(input_variable):
 def channel_names(run_number,verbose=False):
     """Prints channel names for a given run_number or scan object"""
     scan = scan_from_run_number_or_scan(run_number)
-    
+
     channel_list = list(scan[0].keys())
     
     if verbose:
@@ -141,6 +142,30 @@ def process_run(run_number, rois,detector='JF16T03V01', calculate =None, only_sh
                     
     Parallel(n_jobs=n_jobs,verbose=10)(delayed(process_step)(i) for i in range(len(scan)))
 
+def find_offset(run_number, ch1,ch2, offsets=range(-7,7),step=0,verbose = False):
+    """Attempts to find an ideal offset between pids of two channels. Offset is added to channel_2.
+    Returns best offset, best_correlation, [all offsets, all correlations].""""
+    scan = scan_info(run_number)
+    d = scan[step]
+    assert len(d[ch1].shape) ==1, "Channel 2 has more than 1 dimension, can't correlate" 
+    assert len(d[ch2].shape) ==1, "Channel 2 has more than 1 dimension, can't correlate" 
+    corrs = []
+    for offset in offsets:
+        scan = cu.scan_info(run_number)
+        d = scan[step]
+        d[ch2].offset = offset
+        subset = d[ch1,ch2]
+        subset.drop_missing()
+        x = subset[ch1].data       
+        y = subset[ch2].data
+        corrs.append(pearsonr(x,y)[0])
+    best_offset = offsets[np.argmax(corrs)]
+    best_corr   = np.max(corrs)
+    if verbose:
+        print(f'Best correlation for the offset: {best_offset}')
+        print(f'Best correlation value: {best_corr:.2f}')
+    return [best_offset,best_corr,[offsets,corrs]]
+
 class ROI:
     """Definition of region of interest (ROI) in image coordinates. 
        (For images from detector as 2d array with convention bottom == 0th row, left == 0th column).