redone header generation for traces

1 files changed, 58 insertions, 50 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 773077d..691174c 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -1,12 +1,14 @@
 from qolab.file_utils import save_table_with_header
 import numpy as np
+import yaml
 
 class Trace:
     def __init__(self, label):
-        self.label = label
+        self.config = {}
+        self.config['label'] = label
+        self.config['unit'] = None
+        self.config['tags'] = {}
         self.values = None
-        self.unit = None
-        self.tags = {}
 
     def plot(self):
         import matplotlib.pyplot as plt
@@ -14,26 +16,25 @@ class Trace:
         plt.legend()
         plt.grid()
     
+    def getConfig(self):
+        return( self.config )
+
     def getHeader(self, prefix=''):
-        header = []
-        if self.label is not None:
-            header.append(f'label = {self.label}')
-        if self.unit is not None:
-            header.append(f'unit = {self.unit}')
-        for k, v in self.tags.items():
-            header.append(f'{k} = {v}')
+        header = yaml.dump(self.getConfig(), default_flow_style=False, sort_keys=False)
+        header = header.split('\n')
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
 
-    def save(self, fname, item_format='e', header=[]):
+    def save(self, fname, item_format='e', header=[], **kwargs):
         header.extend(self.getHeader())
-        save_table_with_header(fname, self.values, header, item_format=item_format)
+        save_table_with_header(fname, self.values, header, item_format=item_format, **kwargs)
 
 
 class TraceXY:
     def __init__(self, label):
-        self.label = label
-        self.tags = {}
+        self.config = {}
+        self.config['label'] = label
+        self.config['tags'] = {}
         self.x = None
         self.y = None
 
@@ -44,26 +45,30 @@ class TraceXY:
         plt.xlabel(self.x.label)
         plt.grid()
 
+    def getConfig(self):
+        config = self.config.copy()
+        config['TraceX'] = {}
+        config['TraceX'] = self.x.getConfig()
+        config['TraceY'] = {}
+        config['TraceY'] = self.y.getConfig()
+        return( config )
+
     def getHeader(self, prefix=''):
-        header = []
-        if self.label is not None:
-            header.append(f'label = {self.label}')
-        for k, v in self.tags.items():
-            header.append(f'{k} = {v}')
-        header.extend(self.x.getHeader(prefix='x_info: '))
-        header.extend(self.y.getHeader(prefix='y_info: '))
+        header = yaml.dump(self.getConfig(), default_flow_style=False, sort_keys=False)
+        header = header.split('\n')
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
         
-    def save(self, fname, item_format='e', header=[]):
+    def save(self, fname, item_format='e', header=[], **kwargs):
         data=np.concatenate((self.x.values, self.y.values), 1)
         header.extend(self.getHeader())
-        save_table_with_header(fname, data, header, item_format=item_format)
+        save_table_with_header(fname, data, header, item_format=item_format, **kwargs)
 
 class TraceSetSameX:
     def __init__(self, label):
-        self.label = label
-        self.tags = {}
+        self.config = {}
+        self.config['label'] = label
+        self.config['tags'] = {}
         self.x = None
         self.traces={}
 
@@ -71,7 +76,7 @@ class TraceSetSameX:
         if len(self.traces) == 0:
             self.x = tr.x
         trY = tr.y
-        self.traces[tr.label]=trY
+        self.traces[tr.config['label']]=trY
         
     def plot(self):
         import matplotlib.pyplot as plt
@@ -100,24 +105,27 @@ class TraceSetSameX:
         tr.y = self.traces[label]
         return (tr)
     
+    def getConfig(self):
+        config = self.config.copy()
+        config['TraceX'] = {}
+        config['TraceX'] = self.x.getConfig()
+        config['TraceY'] = {}
+        for k,v in self.traces.items():
+            config['TraceY'][k] = v.getConfig()
+        return( config )
+
     def getHeader(self, prefix=''):
-        header = []
-        if self.label is not None:
-            header.append(f'label = {self.label}')
-        for k, v in self.tags.items():
-            header.append(f'{k} = {v}')
-        header.extend(self.x.getHeader(prefix='x_info: '))
-        for k, v in self.traces.items():
-            header.extend(v.getHeader(prefix=f'{k}_info: '))
+        header = yaml.dump(self.getConfig(), default_flow_style=False, sort_keys=False)
+        header = header.split('\n')
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
         
-    def save(self, fname, item_format='e', header=[]):
+    def save(self, fname, item_format='e', header=[], **kwargs):
         data=self.x.values
         for k,v in self.traces.items():
             data=np.concatenate((data, v.values), 1)
         header.extend(self.getHeader())
-        save_table_with_header(fname, data, header, item_format=item_format)
+        save_table_with_header(fname, data, header, item_format=item_format, **kwargs)
 
 
 if __name__ == '__main__': 
@@ -125,30 +133,30 @@ if __name__ == '__main__':
     x=Trace('x trace')
     x.values = np.random.normal(2,2,(4,1))
     x.values = np.array(x.values, int)
-    x.unit='s'
-    x.tags['x_tag'] = 'xxxx'
-    x.tags['x_tag2'] = 'xxxx'
-    x.save('xtrace.dat')
-    print(x.getHeader())
+    x.config['unit']='s'
+    x.config['tags']['tag1'] = 'xxxx'
+    x.config['tags']['tag2'] = 'xxxx'
+    x.save('xtrace.dat', skip_headers_if_file_exist=True)
+    # print(x.getHeader())
     y=Trace('y trace')
     y.values = np.random.normal(2,2,(4,1))
-    y.unit='V'
-    y.tags['y_tag'] = 'yyyy'
+    y.config['unit']='V'
+    y.config['tags']['ytag2'] = 'yyyy'
     xy=TraceXY('xy trace')
-    xy.tags['xy tag']= 'I am xy tag'
+    xy.config['tags']['xy tag']= 'I am xy tag'
     xy.x = x
     xy.y = y
     xy.save('xytrace.dat')
-    print(xy.getHeader())
+    # print(xy.getHeader())
     xyn = TraceSetSameX('many ys trace')
-    xyn.tags['descr'] = 'I am many ys trace'
-    xy.label='y1'
+    xyn.config['tags']['descr'] = 'I am many ys trace'
+    xy.config['label']='y1'
     xyn.addTrace(xy)
-    xy.label='y2'
+    xy.config['label']='y2'
     xyn.addTrace(xy)
-    xy.label='y3'
+    xy.config['label']='y3'
     xyn.addTrace(xy)
     xyn.save('xyntrace.dat')
-    print(xyn.getHeader())
+    # print(xyn.getHeader())