added timestamp type of a plot

1 files changed, 19 insertions, 2 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 7c75d39..009ce2d 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 datetime
 import numpy as np
 import yaml
 
@@ -9,11 +10,19 @@ def headerFromDictionary(d, prefix=''):
     header.extend(tail)
     prefixed_header = [prefix+l for l in header]
     return prefixed_header
+
+def from_timestamps_to_dates(timestamps):
+    dates = [datetime.datetime.fromtimestamp(float(ts)) for ts in timestamps]
+    return(dates)
     
 class Trace:
     def __init__(self, label):
         self.config = {}
         self.config['label'] = label
+        # 'type' is useful to indicate way of representation, make sense on y_vs_x traces
+        # if set to none we have normal y vs x
+        # if set to 'timestamp' x will be converted to datetime dates
+        self.config['type'] = None 
         self.config['tags'] = {}
         self.item_format='.15e'
         self.last_saved_pos = 0
@@ -81,7 +90,11 @@ class TraceXY(Trace):
 
     def plot(self):
         import matplotlib.pyplot as plt
-        plt.plot(self.x.values, self.y.values, label=self.config['label'])
+        x=self.x.values
+        if self.config['type'] is not None:
+            if self.config['type'] == 'timestamp':
+                x = from_timestamps_to_dates(x)
+        plt.plot(x, 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()
@@ -136,13 +149,17 @@ class TraceSetSameX(Trace):
         import matplotlib.pyplot as plt
         nplots = len(self.traces.keys())
         cnt=0
+        x=self.x.values
+        if self.config['type'] is not None:
+            if self.config['type'] == 'timestamp':
+                x = from_timestamps_to_dates(x)
         for k, tr in self.traces.items():
             cnt+=1
             if cnt == 1:
                 ax1=plt.subplot(nplots, 1, cnt)
             else:
                 plt.subplot(nplots, 1, cnt, sharex=ax1)
-            plt.plot(self.x.values, tr.values, label=k)
+            plt.plot(x, tr.values, label=k)
             plt.ylabel(f"{tr.config['label']} ({tr.config['unit']})")
             plt.legend()
             plt.grid()