clear_data method for traces

1 files changed, 20 insertions, 0 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 15711f6..f6c6e68 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -17,7 +17,12 @@ class Trace:
         self.config['unit'] = None
         self.config['tags'] = {}
         self.values = np.empty(0)
+        self.clear_data()
+
+    def clear_data(self):
         self.last_saved_pos = 0
+        if self.values is not None:
+            self.values = np.empty(0, dtype=self.values.dtype)
 
     def plot(self):
         import matplotlib.pyplot as plt
@@ -57,7 +62,14 @@ class TraceXY(Trace):
         self.config['tags'] = {}
         self.x = None
         self.y = None
+        self.clear_data()
+
+    def clear_data(self):
         self.last_saved_pos = 0
+        if self.x is not None:
+            self.x.clear_data()
+        if self.y is not None:
+            self.y.clear_data()
 
     def plot(self):
         import matplotlib.pyplot as plt
@@ -97,7 +109,15 @@ class TraceSetSameX(Trace):
         self.config['tags'] = {}
         self.x = None
         self.traces={}
+        self.clear_data()
+
+    def clear_data(self):
         self.last_saved_pos = 0
+        if self.x is not None:
+            self.x.clear_data()
+        for k, tr in self.traces.items():
+            tr.clear_data()
+
 
     def addTrace(self, tr):
         if len(self.traces) == 0: