From 66930122523e1993f4aed64f7401fede8931e32e Mon Sep 17 00:00:00 2001
From: Vonka Jakub <jakub.vonka@psi.ch>
Date: Wed, 14 Sep 2022 12:41:36 +0200
Subject: [PATCH] Added process_run() to create small processed data and
 scan_info() to read all channels including the small data ones.

---
 src/cristallina/utils.py | 81 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 76 insertions(+), 5 deletions(-)

diff --git a/src/cristallina/utils.py b/src/cristallina/utils.py
index 668de38..508b912 100644
--- a/src/cristallina/utils.py
+++ b/src/cristallina/utils.py
@@ -1,9 +1,25 @@
 import yaml
 import os
 import numpy as np
-from sfdata import SFDataFiles, sfdatafile, SFScanInfo
+from sfdata import SFDataFiles, sfdatafile, SFScanInfo, SFProcFile
 from xraydb import material_mu
+from joblib import Parallel, delayed
 
+def scan_info(run_number,base_path=None,small_data=True):
+    """Returns SFScanInfo object for a given run number. 
+    If there is are small data channels, they will be added (small_data=False to suppress their loading).
+    """
+    if base_path == None:
+        base_path = heuristic_extract_base_path()
+    
+    scan = SFScanInfo(f"{base_path}/run{run_number:04}/meta/scan.json")
+    
+    if small_data:
+        for i in range(len(scan.readbacks)):
+            sd_path_for_step=heuristic_extract_smalldata_path()+'run'+str(run_number).zfill(4)+'/acq'+str(i+1).zfill(4)+'.smalldata.h5'
+            scan.info['scan_files'][i].append(sd_path_for_step)
+
+    return scan
 
 def print_run_info(
     run_number=42, print_channels=True, extra_verbose=False, base_path=None,
@@ -12,10 +28,10 @@ def print_run_info(
 
     Extra verbose output contains all files and pids.
     """
-    if base_path == None:
-    	base_path = heuristic_extract_base_path()
-    scan = SFScanInfo(f"{base_path}/run{run_number:04}/meta/scan.json")
-
+    # if base_path == None:
+    # 	base_path = heuristic_extract_base_path()
+    # scan = SFScanInfo(f"{base_path}/run{run_number:04}/meta/scan.json")
+    scan = scan_info(run_number,base_path=base_path)
     short = {}
     for key, value in scan.info.items():
         if isinstance(value, list):
@@ -49,6 +65,61 @@ def print_run_info(
         break
 
 
+def process_run(run_number, rois,detector='JF16T03V01', roi_img=True, calculate =None, only_shots=slice(None), n_jobs=12):
+    """Process rois for a given detector. Save the results small data in the res/small_data/run...
+    By default only sum of rois is calculated, [mean,std,img] can be added to the "calculate" optional parameter.
+    """    
+    # Load scan object with SFScanInfo
+    scan = scan_info(run_number,small_data=False)
+    
+    # Set the path for later small data saving
+    path_with_run_folder = heuristic_extract_smalldata_path()+'run'+str(run_number).zfill(4)
+    
+    # Make the small data run folder if it doesn't exist
+    if not os.path.exists( path_with_run_folder ):
+        os.mkdir( path_with_run_folder )
+    
+    # Add scan info into the small_data_object
+    #for key in scan.info.keys():
+    #    d[key]= info.info[key]
+    
+    # Check if there is only one roi. If yes, make a list of it so it can be iterated over.
+    if isinstance(rois, ROI):
+        rois = [rois]
+
+    def process_step(i):
+        scan = scan_info(run_number,small_data=False)
+        
+        step = scan[i]
+        with step as data:
+            with SFProcFile(f"{path_with_run_folder}/acq{str(i+1).zfill(4)}.smalldata.h5", mode="w") as sd:
+                subset = data[detector]
+
+                # Calculate everything related to JF_rois
+                for roi in rois:
+
+                    bottom, top, left, right = roi.bottom, roi.top, roi.left, roi.right
+                    roi_name = roi.name
+
+                    # Pulse ids for saving the new channels
+                    det_pids = data[detector].pids
+                    sd[roi.name] = (det_pids[only_shots], data[detector][only_shots, left:right, bottom:top].sum(axis=(1, 2)))
+                    if calculate:
+                        if 'mean' in calculate:
+                            sd[roi.name+"_mean"] = (det_pids[only_shots], data[detector][only_shots, left:right, bottom:top].mean(axis=(1, 2)))
+                        if 'std' in calculate:
+                            sd[roi.name+"_std"] = (det_pids[only_shots], data[detector][only_shots, left:right, bottom:top].std(axis=(1, 2)))
+                        if 'img' in calculate:
+                            sd[f'{roi.name}_img'] = (det_pids[only_shots], data[detector][only_shots, left:right, bottom:top].data)
+                    
+                    # Currently meta files can't be read by SFData, this will be modified by Sven and then we can use it. For now saving in roi_info
+                    #sd.meta[roi.name+"_info"] = f"roi {roi.name}: {left},{right}; {bottom},{top} (left, right, bottom, top)"
+
+                    # These channels have only one dataset per step of the scan, so we take the first pulseID
+                    sd[roi.name + "_info"]   =([det_pids[0]], [f"roi {roi.name}: {left},{right}; {bottom},{top} (left, right, bottom, top)"])
+                    sd[roi.name + "_mean_img"] = ([det_pids[0]], [data[detector][:, left:right, bottom:top].mean(axis=(0))] )                    
+                    
+    Parallel(n_jobs=n_jobs,verbose=10)(delayed(process_step)(i) for i in range(len(scan)))
 
 class ROI:
     """Definition of region of interest (ROI) in image coordinates.