added grabing and ploting all traces capabilities

1 files changed, 40 insertions, 5 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 9b14264..ef44041 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -1,16 +1,51 @@
-
-class Trace:
+class TraceYonly:
     def __init__(self, descrStr):
         self.descr = descrStr
-        self.x = None
-        self.xlabel = None
-        self.xunit = None
         self.y = None
         self.ylabel = None
         self.yunit = None
 
     def plot(self):
         import matplotlib.pyplot as plt
+        plt.plot(self.y, label=self.descr)
+        plt.legend()
+        
+
+class Trace(TraceYonly):
+    def __init__(self, descrStr):
+        super().__init__(descrStr)
+        self.x = None
+        self.xlabel = None
+        self.xunit = None
+
+    def plot(self):
+        import matplotlib.pyplot as plt
         plt.plot(self.x, self.y, label=self.descr)
         plt.legend()
         
+class TraceSetSameX:
+    def __init__(self, descrStr):
+        self.descr = descrStr
+        self.x = None
+        self.xlabel = None
+        self.xunit = None
+        self.traces={}
+
+    def addTrace(self, tr):
+        if len(self.traces) == 0:
+            self.x = tr.x
+            self.xlabel = tr.xlabel
+            self.xunit  = tr.xunit
+        trY = TraceYonly(tr.descr)
+        trY.y = tr.y
+        trY.ylabel = tr.ylabel
+        trY.yunit  = tr.yunit
+        self.traces[trY.descr]=trY
+        
+    def plot(self):
+        import matplotlib.pyplot as plt
+        for k, v in self.traces.items():
+            plt.plot(self.x, v.y, label=v.descr)
+        plt.legend()
+
+