From 3c435f6f9be84f15dbdfa43ad9794f81696f3367 Mon Sep 17 00:00:00 2001
From: vonka_j <jakub.vonka@psi.ch>
Date: Mon, 31 Oct 2022 22:04:26 +0000
Subject: [PATCH] update find_offset()

---
 src/cristallina/utils.py | 34 ++++++++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/cristallina/utils.py b/src/cristallina/utils.py
index 84666d3..8f1c411 100644
--- a/src/cristallina/utils.py
+++ b/src/cristallina/utils.py
@@ -142,29 +142,51 @@ 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.
+def find_offset(run_no_or_scan, ch1,ch2, offsets=range(-7,7),step=0,verbose = False, plot = False):
+    """Attempts to find an ideal offset between pids of two channels. Offset is added to Channel_2. Channel_1 is used as reference.
     Returns best offset, best_correlation, [all offsets, all correlations]."""
-    scan = scan_info(run_number)
+    scan = scan_from_run_number_or_scan(run_no_or_scan)
+    run_number = int(str(scan.fs)[-19:-15])
     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 = []
+    corrs,xs,ys = [],[],[]
+    
     for offset in offsets:
-        scan = scan_info(run_number)
+        scan = scan_from_run_number_or_scan(run_no_or_scan)
         d = scan[step]
         d[ch2].offset = offset
         subset = d[ch1,ch2]
         subset.drop_missing()
         x = subset[ch1].data       
         y = subset[ch2].data
+        xs.append(x)
+        ys.append(y)
         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]]
+    if plot:
+        fig,ax  = plt.subplots(1,2,figsize=(12,4))
+        for i,offset in enumerate(offsets):
+            if offset == best_offset:
+                col = 'tab:red'
+                alpha = 1
+            else:
+                col = 'tab:grey'
+                alpha =0.2
+            ax[0].scatter(xs[i],ys[i],color=col,alpha = alpha )
+        fig.suptitle(f'Run {run_number}, Best offset {best_offset}, Correlation: {best_corr:.3f}')
+        ax[0].set_xlabel(ch1)
+        ax[0].set_ylabel(ch2)
+        ax[1].plot(offsets,corrs,'-o',color='tab:blue')
+        ax[1].set_xlabel('Offset')
+        ax[1].set_ylabel('Correlation')
+        ax[1].plot(best_offset,best_corr,'o',markersize=13,color='tab:red',mfc='none')
+    return best_offset,best_corr,[offsets,corrs]
+
 
 class ROI:
     """Definition of region of interest (ROI) in image coordinates.