diff options
| -rw-r--r-- | qolab/data/trace.py | 64 |
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__': |