new class ess.TraceWin.envelope

ess/TraceWin.py

@@ -1265,6 +1265,99 @@ class envDiag:
         return ret
 
 
+class envelope:
+    """
+    Read an envelope file
+    Create one by saving envelope data plot
+    to ascii
+
+    Example::
+        from ess import TraceWin
+        from matplotlib import pyplot as plt
+        data = TraceWin.envelope('envelope.txt')
+        print(data.keys())
+        for key in data:
+            print(key, data.unit(key))
+            if 'rms_' in key:
+                plt.plot(data['position'], data[key]/max(data[key]), label=f"{key} [{data.unit(key)}]")
+        plt.legend()
+        plt.xlabel(f"Position [{data.unit('position')}]")
+        plt.show()
+
+    """
+
+    def __init__(self, filename):
+        self.filename = filename
+        self.headers = ()
+        self._units = []
+        self._raw_data = None
+        self._readAsciiFile()
+
+    def _readAsciiFile(self):
+
+        import numpy
+
+        self._raw_data = numpy.loadtxt(self.filename, skiprows=1)
+        with open(self.filename, "r") as fin:
+            header = fin.readline()
+            print(header)
+            headers = []
+            for h in header.split():
+                if "centroid" == h:
+                    continue
+                elif "(" not in h:
+                    headers.append(h)
+                elif "unit(" in h:
+                    self._units = tuple(h.split("(")[1][:-1].split(","))
+                else:
+                    base, main = h.split("(")
+                    if base:
+                        for k in main[:-1].split(","):
+                            headers.append(f"{base}_{k}")
+                    else:
+                        headers.extend(main[:-1].split(","))
+            self.headers = tuple(headers)
+
+    def __iter__(self):
+        return iter(self.keys())
+
+    def keys(self):
+        return self.headers
+
+    def where(self, column):
+        if column not in self.keys():
+            raise ValueError(f"Wrong column name {column}")
+        for i in range(len(self.keys())):
+            if self.keys()[i] == column:
+                return i
+
+    def __getitem__(self, column):
+        """
+        Get the data of the column specified
+
+        """
+        index = self.where(column)
+        return self._raw_data[:, index]
+
+    def unit(self, column):
+        """
+        TODO gam-1
+        """
+        if "'" in column:
+            return self._units[1]
+        elif "/" in column:
+            return ""
+        elif column in ["time"]:
+            return self._units[3]
+        elif "phase" in column:
+            return self._units[2]
+        elif "energy" in column:
+            return self._units[4]
+        elif "gam-1" == column:
+            return "GeV"
+        return self._units[0]
+
+
 class partran(dict):
     """
     Read partran1.out files..