added write for sameX trace

1 files changed, 16 insertions, 3 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index dccc2f7..783a649 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -1,4 +1,5 @@
 from qolab.file_utils import save_table_with_header
+import numpy as np
 
 class Trace:
     def __init__(self, label):
@@ -51,6 +52,10 @@ class TraceXY:
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
         
+    def write(self, fname, item_format='e'):
+        data=np.concatenate((self.x.values, self.y.values), 1)
+        save_table_with_header(fname, data, self.header(), item_format=item_format)
+
 class TraceSetSameX:
     def __init__(self, label):
         self.label = label
@@ -95,24 +100,31 @@ class TraceSetSameX:
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
         
+    def write(self, fname, item_format='e'):
+        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)
+
 
 if __name__ == '__main__': 
-    import numpy as np
     print("Testing trace")
     x=Trace('x trace')
-    x.values = np.random.normal(2,2,(4,2))
+    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.write('xtrace.dat')
     y=Trace('y trace')
-    y.values = np.random.normal(2,2,(4,2))
+    y.values = np.random.normal(2,2,(4,1))
     y.unit='V'
     y.tags['y_tag'] = 'yyyy'
     xy=TraceXY('xy trace')
     xy.tags['xy tag']= 'I am xy tag'
     xy.x = x
     xy.y = y
+    xy.write('xytrace.dat')
     print(xy.header())
     xyn = TraceSetSameX('many ys trace')
     xyn.tags['descr'] = 'I am many ys trace'
@@ -122,6 +134,7 @@ if __name__ == '__main__':
     xyn.addTrace(xy)
     xy.label='y3'
     xyn.addTrace(xy)
+    xyn.write('xyntrace.dat')
     print(xyn.header())