Refactor split_scans

nafuma/xanes/io.py

@@ -3,93 +3,119 @@ import matplotlib.pyplot as plt
 import os
 import numpy as np
 import nafuma.auxillary as aux
+from nafuma.xanes.calib import find_element
 
 
-def split_xanes_scan(root, destination=None, replace=False):
+def split_scan_data(data: dict, options={}):
+    
+
+    required_options = ['save', 'save_folder', 'replace', 'add_rois']
+
+    default_options = {
+        'save': False,
+        'save_folder': '.',
+        'replace': False,
+        'add_rois': False
+    }
+
+    options = aux.update_options(options=options, required_options=required_options, default_options=default_options)
     #root is the path to the beamtime-folder
     #destination should be the path to the processed data
     
     #insert a for-loop to go through all the folders.dat-files in the folder root\xanes\raw
 
     # FIXME Only adding this variable to pass the Linting-tests - will refactor this later
-    filename = 'dummy'
+
+    if not isinstance(data['path'], list):
+        data['path'] = [data['path']]
+
+    all_scans = []
+    
+    for filename in data['path']:
 
         with open(filename, 'r') as f:
             lines = f.readlines()
             
-    datas = []
-    data = []
-    headers = []
-    header = ''
-    start = False
+        scan_datas, scan_data = [], []
+        headers, header = [], ''
+        read_data = False
         
         for line in lines:
+            # Header line starts with #L - reads headers, and toggles data read-in on
             if line[0:2] == "#L":
-            start = True
-            header = line[2:].split()
+                header, read_data = line[2:].split(), True
                 continue
 
+            # First line after data started with #C - stops data read-in
             elif line[0:2] == "#C":
-            start = False
+                read_data = False
                 
-            if data:
-                datas.append(data)
-                data = []
+                if scan_data:
+                    scan_datas.append(scan_data); scan_data = []
                     
                 if header:
-                headers.append(header)
-                header = ''
+                    headers.append(header); header = ''
                     
-                
-
-        if start == False:
+            # Ignore line if read-in not toggled       
+            if read_data == False:
                 continue
             
+            # Read in data if it is
             else:
-            data.append(line.split())
+                scan_data.append(line.split())
                 
                 
+        edges = {'Mn': [], 'Fe': [], 'Co': [], 'Ni': []}
         
 
-    edges = {'Mn': [6.0, 6.1, 6.2, 6.3, 6.4, 6.5], 'Fe': [6.8, 6.9, 7.0, 7.1, 7.2], 'Co': [7.6, 7.7, 7.8, 7.9], 'Ni': [8.1, 8.2, 8.3, 8.4, 8.5]}
-    edge_count = {'Mn': 0, 'Fe': 0, 'Co': 0, 'Ni': 0}
+        for i, scan_data in enumerate(scan_datas):
+            xanes_df = pd.DataFrame(scan_data).apply(pd.to_numeric)
+            xanes_df.columns = headers[i]
+
+            if not ('xmap_roi00' in headers[i]) and (not 'xmap_roi01' in headers[i]):
+                continue 
 
             
-    for ind, data in enumerate(datas):
-        df = pd.DataFrame(data)
-        df.columns = headers[ind]
-
-        edge_start = np.round((float(df["ZapEnergy"].min())), 1)
-
-        for edge, energies in edges.items():
-            if edge_start in energies:
-                edge_actual = edge
-                edge_count[edge] += 1
+            edge = find_element({'xanes_data_original': xanes_df})
+            edges[edge].append(xanes_df)
             
         
+        if options['add']:
    
-        filename = filename.split('/')[-1]
-        count = str(edge_count[edge_actual]).zfill(4)
+            added_edges = {'Mn': [], 'Fe': [], 'Co': [], 'Ni': []}
+            for edge, scans in edges.items():
+                if scans:
+                    xanes_df = scans[0]
+
+                    for i, scan in enumerate(scans):
+                        if i > 0:
+
+                            if 'xmap_roi00' in xanes_df.columns:
+                                xanes_df['xmap_roi00'] += scan['xmap_roi00']
+                            if 'xmap_roi01' in xanes_df.columns:
+                                xanes_df['xmap_roi01'] += scan['xmap_roi01']
+
+                    added_edges[edge].append(xanes_df)
+
+            edges = added_edges
+            
+        if options['save']:
+            if not os.path.isdir(options['save_folder']):
+                os.makedirs(options['save_folder'])
 
 
-        # Save 
-        if destination:
-            cwd = os.getcwd()
+            filename = os.path.basename(filename).split('.')[0]
 
-            if not os.path.isdir(destination):
-                os.mkdir(destination)
-                
-            os.chdir(destination)
-
-            df.to_csv('{}_{}_{}.dat'.format(filename.split('.')[0], edge_actual, count))
-
-            os.chdir(cwd)
-
-        else:
-            df.to_csv('{}_{}_{}.dat'.format(filename.split('.')[0], edge_actual, count))
+            for edge, scans in edges.items():
+                for i, scan in enumerate(scans):
+                    count = '' if options['add'] else '_'+str(i).zfill(4)
+                    path = os.path.join(options['save_folder'], f'{filename}_{edge}{count}.dat')
+                    scan.to_csv(path)
 
+        all_scans.append(edges)
 
 
+    return all_scans
 
 
 
@@ -117,6 +143,7 @@ def read_data(data: dict, options={}) -> pd.DataFrame:
         columns.append(filename)
 
         scan_data = pd.read_csv(filename)
+
         scan_data = scan_data[[determine_active_roi(scan_data)]]
         xanes_data = pd.concat([xanes_data, scan_data], axis=1)