traces take external header argument

1 files changed, 22 insertions, 18 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 8478447..773077d 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -14,7 +14,7 @@ class Trace:
         plt.legend()
         plt.grid()
     
-    def header(self, prefix=''):
+    def getHeader(self, prefix=''):
         header = []
         if self.label is not None:
             header.append(f'label = {self.label}')
@@ -25,8 +25,9 @@ class Trace:
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
 
-    def save(self, fname, item_format='e'):
-        save_table_with_header(fname, self.values, self.header(), item_format=item_format)
+    def save(self, fname, item_format='e', header=[]):
+        header.extend(self.getHeader())
+        save_table_with_header(fname, self.values, header, item_format=item_format)
 
 
 class TraceXY:
@@ -43,20 +44,21 @@ class TraceXY:
         plt.xlabel(self.x.label)
         plt.grid()
 
-    def header(self, prefix=''):
+    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.header(prefix='x_info: '))
-        header.extend(self.y.header(prefix='y_info: '))
+        header.extend(self.x.getHeader(prefix='x_info: '))
+        header.extend(self.y.getHeader(prefix='y_info: '))
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
         
-    def save(self, fname, item_format='e'):
+    def save(self, fname, item_format='e', header=[]):
         data=np.concatenate((self.x.values, self.y.values), 1)
-        save_table_with_header(fname, data, self.header(), item_format=item_format)
+        header.extend(self.getHeader())
+        save_table_with_header(fname, data, header, item_format=item_format)
 
 class TraceSetSameX:
     def __init__(self, label):
@@ -98,23 +100,24 @@ class TraceSetSameX:
         tr.y = self.traces[label]
         return (tr)
     
-    def header(self, prefix=''):
+    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.header(prefix='x_info: '))
+        header.extend(self.x.getHeader(prefix='x_info: '))
         for k, v in self.traces.items():
-            header.extend(v.header(prefix=f'{k}_info: '))
+            header.extend(v.getHeader(prefix=f'{k}_info: '))
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
         
-    def save(self, fname, item_format='e'):
+    def save(self, fname, item_format='e', header=[]):
         data=self.x.values
         for k,v in self.traces.items():
             data=np.concatenate((data, v.values), 1)
-        save_table_with_header(fname, data, self.header(), item_format=item_format)
+        header.extend(self.getHeader())
+        save_table_with_header(fname, data, header, item_format=item_format)
 
 
 if __name__ == '__main__': 
@@ -125,7 +128,8 @@ if __name__ == '__main__':
     x.unit='s'
     x.tags['x_tag'] = 'xxxx'
     x.tags['x_tag2'] = 'xxxx'
-    x.write('xtrace.dat')
+    x.save('xtrace.dat')
+    print(x.getHeader())
     y=Trace('y trace')
     y.values = np.random.normal(2,2,(4,1))
     y.unit='V'
@@ -134,8 +138,8 @@ if __name__ == '__main__':
     xy.tags['xy tag']= 'I am xy tag'
     xy.x = x
     xy.y = y
-    xy.write('xytrace.dat')
-    print(xy.header())
+    xy.save('xytrace.dat')
+    print(xy.getHeader())
     xyn = TraceSetSameX('many ys trace')
     xyn.tags['descr'] = 'I am many ys trace'
     xy.label='y1'
@@ -144,7 +148,7 @@ if __name__ == '__main__':
     xyn.addTrace(xy)
     xy.label='y3'
     xyn.addTrace(xy)
-    xyn.write('xyntrace.dat')
-    print(xyn.header())
+    xyn.save('xyntrace.dat')
+    print(xyn.getHeader())