added header generation for traces

2 files changed, 55 insertions, 3 deletions

diff --git a/qolab/data/__init__.py b/qolab/data/__init__.py
index 528815a..c0f071c 100644
--- a/qolab/data/__init__.py
+++ b/qolab/data/__init__.py
@@ -1,6 +1,8 @@
 
-from .trace import Trace
+from .trace import Trace, TraceXY, TraceSetSameX
 
 __all__ = [
     "Trace",
+    "TraceXY",
+    "TraceSetSameX",
 ]
diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index bf7e54f..562224d 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -11,7 +11,7 @@ class Trace:
         plt.plot(self.values, label=self.descr)
         plt.legend()
     
-    def header(self):
+    def header(self, prefix=''):
         header = []
         if self.label is not None:
             header.append(f'label = {self.label}')
@@ -19,7 +19,8 @@ class Trace:
             header.append(f'unit = {self.unit}')
         for k, v in self.tags.items():
             header.append(f'{k} = {v}')
-        return header
+        prefixed_header = [prefix+l for l in header]
+        return prefixed_header
 
     def write(self, fname):
         pass
@@ -29,6 +30,7 @@ class Trace:
 class TraceXY:
     def __init__(self, label):
         self.label = label
+        self.tags = {}
         self.x = None
         self.y = None
 
@@ -37,10 +39,22 @@ class TraceXY:
         plt.plot(self.x.values, self.y.values, label=self.label)
         plt.legend()
         plt.xlabel(self.x.label)
+
+    def header(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_'))
+        prefixed_header = [prefix+l for l in header]
+        return prefixed_header
         
 class TraceSetSameX:
     def __init__(self, label):
         self.label = label
+        self.tags = {}
         self.x = None
         self.traces={}
 
@@ -69,3 +83,39 @@ class TraceSetSameX:
         tr.y = self.traces[label]
         return (tr)
     
+    def header(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_'))
+        for k, v in self.traces.items():
+            header.extend(v.header(prefix=f'{k}_'))
+        prefixed_header = [prefix+l for l in header]
+        return prefixed_header
+        
+
+if __name__ == '__main__': 
+    x=Trace('x trace')
+    x.unit='s'
+    x.tags['x_tag'] = 'xxxx'
+    x.tags['x_tag2'] = 'xxxx'
+    y=Trace('y trace')
+    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
+    print(xy.header())
+    xyn = TraceSetSameX('many ys trace')
+    xyn.tags['descr'] = 'I am many ys trace'
+    xyn.addTrace(xy)
+    xy.label='y2'
+    xyn.addTrace(xy)
+    xy.label='y3'
+    xyn.addTrace(xy)
+    print(xyn.header())
+
+