simplified trace save and plot

1 files changed, 31 insertions, 33 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 691174c..107caae 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -2,6 +2,14 @@ from qolab.file_utils import save_table_with_header
 import numpy as np
 import yaml
 
+def headerFronDictionary(d, prefix=''):
+    header = []
+    tail = yaml.dump(d, default_flow_style=False, sort_keys=False)
+    tail = tail.split('\n')
+    header.extend(tail)
+    prefixed_header = [prefix+l for l in header]
+    return prefixed_header
+    
 class Trace:
     def __init__(self, label):
         self.config = {}
@@ -12,25 +20,26 @@ class Trace:
 
     def plot(self):
         import matplotlib.pyplot as plt
-        plt.plot(self.values, label=self.descr)
+        plt.plot(self.values, label=self.config['label'])
+        plt.xlabel('index')
+        plt.ylabel(f"{self.config['unit']}")
         plt.legend()
         plt.grid()
     
     def getConfig(self):
         return( self.config )
 
+    def getData(self):
+        return( self.values )
+
     def getHeader(self, prefix=''):
-        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
+        return headerFronDictionary(self.getConfig(), prefix='')
 
-    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, **kwargs)
+    def save(self, fname, item_format='e', **kwargs):
+        save_table_with_header(fname, self.getData(), self.getHeader(), item_format=item_format, **kwargs)
 
 
-class TraceXY:
+class TraceXY(Trace):
     def __init__(self, label):
         self.config = {}
         self.config['label'] = label
@@ -40,9 +49,10 @@ class TraceXY:
 
     def plot(self):
         import matplotlib.pyplot as plt
-        plt.plot(self.x.values, self.y.values, label=self.label)
+        plt.plot(self.x.values, self.y.values, label=self.config['label'])
+        plt.xlabel(f"{self.x.config['label']} ({self.x.config['unit']})")
+        plt.ylabel(f"{self.y.config['label']} ({self.y.config['unit']})")
         plt.legend()
-        plt.xlabel(self.x.label)
         plt.grid()
 
     def getConfig(self):
@@ -53,18 +63,11 @@ class TraceXY:
         config['TraceY'] = self.y.getConfig()
         return( config )
 
-    def getHeader(self, prefix=''):
-        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=[], **kwargs):
+    def getData(self):
         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, **kwargs)
-
-class TraceSetSameX:
+        return( data )
+        
+class TraceSetSameX(Trace):
     def __init__(self, label):
         self.config = {}
         self.config['label'] = label
@@ -88,10 +91,12 @@ class TraceSetSameX:
                 ax1=plt.subplot(nplots, 1, cnt)
             else:
                 plt.subplot(nplots, 1, cnt, sharex=ax1)
-            plt.plot(self.x.values, tr.values, label=tr.label)
-            plt.xlabel(self.x.label)
+            plt.plot(self.x.values, tr.values, label=k)
+            plt.xlabel(f"{self.x.config['label']} ({self.x.config['unit']})")
+            plt.ylabel(f"{tr.config['label']} ({tr.config['unit']})")
             plt.legend()
             plt.grid()
+            # plt.plot(self.x.values, tr.values, label=tr.label)
 
     def items(self):
         return (self.traces.items())
@@ -114,18 +119,11 @@ class TraceSetSameX:
             config['TraceY'][k] = v.getConfig()
         return( config )
 
-    def getHeader(self, prefix=''):
-        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=[], **kwargs):
+    def getData(self):
         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, **kwargs)
+        return( data )
 
 
 if __name__ == '__main__':