From fe1d8044a3f2f6e24927a4f8baf4aa0894a56140 Mon Sep 17 00:00:00 2001
From: Ivan Usov <ivan.usov@psi.ch>
Date: Tue, 15 Sep 2020 16:07:41 +0200
Subject: [PATCH] Simplify dat file processing

* rename "omega" -> "om", and "counts" -> "Counts" for consistency
with the column names
---
 pyzebra/ccl_findpeaks.py |  4 ++--
 pyzebra/fit2.py          |  4 ++--
 pyzebra/load_1D.py       | 44 ++++++++++++++++++----------------------
 3 files changed, 24 insertions(+), 28 deletions(-)

diff --git a/pyzebra/ccl_findpeaks.py b/pyzebra/ccl_findpeaks.py
index fc52425..e141f95 100644
--- a/pyzebra/ccl_findpeaks.py
+++ b/pyzebra/ccl_findpeaks.py
@@ -58,8 +58,8 @@ def ccl_findpeaks(
         print("Invalid value for prominence, select positive number, new value set to:", prominence)
 
 
-    omega = data["Measurements"][str(keys)]["omega"]
-    counts = np.array(data["Measurements"][str(keys)]["counts"])
+    omega = data["Measurements"][str(keys)]["om"]
+    counts = np.array(data["Measurements"][str(keys)]["Counts"])
     if smooth is True:
         itp = interp1d(omega, counts, kind="linear")
         absintensity = [abs(number) for number in counts]
diff --git a/pyzebra/fit2.py b/pyzebra/fit2.py
index ac499e3..be5fc22 100644
--- a/pyzebra/fit2.py
+++ b/pyzebra/fit2.py
@@ -40,8 +40,8 @@ def fitccl(
         print("NO PEAK or more than 1 peak")
         return
 
-    x = list(data["Measurements"][str(keys)]["omega"])
-    y = list(data["Measurements"][str(keys)]["counts"])
+    x = list(data["Measurements"][str(keys)]["om"])
+    y = list(data["Measurements"][str(keys)]["Counts"])
     peak_index = data["Measurements"][str(keys)]["peak_indexes"]
     peak_height = data["Measurements"][str(keys)]["peak_heights"]
     print("before", constraints_min)
diff --git a/pyzebra/load_1D.py b/pyzebra/load_1D.py
index caf96ee..30ff5b8 100644
--- a/pyzebra/load_1D.py
+++ b/pyzebra/load_1D.py
@@ -1,7 +1,9 @@
 import re
-import numpy as np
+from collections import defaultdict
 from decimal import Decimal
 
+import numpy as np
+
 META_VARS_STR = (
     "instrument",
     "title",
@@ -133,14 +135,14 @@ def load_1D(filepath):
                         fileline = data[position + 2 + i].split()
                         numbers = [int(w) for w in fileline]
                         counts = counts + numbers
-                    d["omega"] = np.linspace(
+                    d["om"] = np.linspace(
                         float(lines.split()[5])
                         - (int(next_line.split()[0]) / 2) * float(next_line.split()[1]),
                         float(lines.split()[5])
                         + (int(next_line.split()[0]) / 2) * float(next_line.split()[1]),
                         int(next_line.split()[0]),
                     )
-                    d["counts"] = counts
+                    d["Counts"] = counts
                     det_variables["Measurements"][str("M" + str(measurement_number))] = d
 
                     if all(decimal):
@@ -149,27 +151,21 @@ def load_1D(filepath):
                         det_variables["meta"]["indices"] = "real"
 
         elif det_variables["file_type"] == "dat":
-            data = infile.readlines()
-            num_of_points = int(data[1].split()[0])
-            omega = []
-            counts = []
-            monitor1 = []
-            monitor2 = []
-            monitor3 = []
-            time = []
-            for position in range(num_of_points):
-                omega.append(float(data[position + 3].split()[1]))
-                counts.append(float(data[position + 3].split()[2]))
-                monitor1.append(float(data[position + 3].split()[3]))
-                monitor2.append(float(data[position + 3].split()[4]))
-                monitor3.append(float(data[position + 3].split()[5]))
-                time.append(float(data[position + 3].split()[6]))
-            det_variables["Measurements"]["omega"] = omega
-            det_variables["Measurements"]["counts"] = counts
-            det_variables["Measurements"]["Monitor1"] = monitor1
-            det_variables["Measurements"]["Monitor2"] = monitor2
-            det_variables["Measurements"]["Monitor3"] = monitor3
-            det_variables["Measurements"]["time"] = time
+            # skip the first 2 rows, the third row contans the column names
+            next(infile)
+            next(infile)
+            row_names = next(infile).split()
+
+            data_cols = defaultdict(list)
+            for line in infile:
+                if "END-OF-DATA" in line:
+                    # this is the end of data
+                    break
+
+                for name, val in zip(row_names, line.split()):
+                    data_cols[name].append(float(val))
+
+            det_variables["Measurements"] = dict(data_cols)
         else:
             print("Unknown file extention")