diff options
Diffstat (limited to 'qolab/data')
| -rw-r--r-- | qolab/data/trace.py | 108 |
1 files changed, 58 insertions, 50 deletions
diff --git a/qolab/data/trace.py b/qolab/data/trace.py index 773077d..691174c 100644 --- a/qolab/data/trace.py +++ b/qolab/data/trace.py @@ -1,12 +1,14 @@ from qolab.file_utils import save_table_with_header import numpy as np +import yaml class Trace: def __init__(self, label): - self.label = label + self.config = {} + self.config['label'] = label + self.config['unit'] = None + self.config['tags'] = {} self.values = None - self.unit = None - self.tags = {} def plot(self): import matplotlib.pyplot as plt @@ -14,26 +16,25 @@ class Trace: plt.legend() plt.grid() + def getConfig(self): + return( self.config ) + def getHeader(self, prefix=''): - header = [] - if self.label is not None: - header.append(f'label = {self.label}') - if self.unit is not None: - header.append(f'unit = {self.unit}') - for k, v in self.tags.items(): - header.append(f'{k} = {v}') + 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=[]): + 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) + save_table_with_header(fname, self.values, header, item_format=item_format, **kwargs) class TraceXY: def __init__(self, label): - self.label = label - self.tags = {} + self.config = {} + self.config['label'] = label + self.config['tags'] = {} self.x = None self.y = None @@ -44,26 +45,30 @@ class TraceXY: plt.xlabel(self.x.label) plt.grid() + def getConfig(self): + config = self.config.copy() + config['TraceX'] = {} + config['TraceX'] = self.x.getConfig() + config['TraceY'] = {} + config['TraceY'] = self.y.getConfig() + return( config ) + def getHeader(self, prefix=''): - header = [] - if self.label is not None: - header.append(f'label = {self.label}') - for k, v in self.tags.items(): - header.append(f'{k} = {v}') - header.extend(self.x.getHeader(prefix='x_info: ')) - header.extend(self.y.getHeader(prefix='y_info: ')) + 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=[]): + def save(self, fname, item_format='e', header=[], **kwargs): 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) + save_table_with_header(fname, data, header, item_format=item_format, **kwargs) class TraceSetSameX: def __init__(self, label): - self.label = label - self.tags = {} + self.config = {} + self.config['label'] = label + self.config['tags'] = {} self.x = None self.traces={} @@ -71,7 +76,7 @@ class TraceSetSameX: if len(self.traces) == 0: self.x = tr.x trY = tr.y - self.traces[tr.label]=trY + self.traces[tr.config['label']]=trY def plot(self): import matplotlib.pyplot as plt @@ -100,24 +105,27 @@ class TraceSetSameX: tr.y = self.traces[label] return (tr) + def getConfig(self): + config = self.config.copy() + config['TraceX'] = {} + config['TraceX'] = self.x.getConfig() + config['TraceY'] = {} + for k,v in self.traces.items(): + config['TraceY'][k] = v.getConfig() + return( config ) + def getHeader(self, prefix=''): - header = [] - if self.label is not None: - header.append(f'label = {self.label}') - for k, v in self.tags.items(): - header.append(f'{k} = {v}') - header.extend(self.x.getHeader(prefix='x_info: ')) - for k, v in self.traces.items(): - header.extend(v.getHeader(prefix=f'{k}_info: ')) + 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=[]): + def save(self, fname, item_format='e', header=[], **kwargs): 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) + save_table_with_header(fname, data, header, item_format=item_format, **kwargs) if __name__ == '__main__': @@ -125,30 +133,30 @@ if __name__ == '__main__': x=Trace('x trace') x.values = np.random.normal(2,2,(4,1)) x.values = np.array(x.values, int) - x.unit='s' - x.tags['x_tag'] = 'xxxx' - x.tags['x_tag2'] = 'xxxx' - x.save('xtrace.dat') - print(x.getHeader()) + x.config['unit']='s' + x.config['tags']['tag1'] = 'xxxx' + x.config['tags']['tag2'] = 'xxxx' + x.save('xtrace.dat', skip_headers_if_file_exist=True) + # print(x.getHeader()) y=Trace('y trace') y.values = np.random.normal(2,2,(4,1)) - y.unit='V' - y.tags['y_tag'] = 'yyyy' + y.config['unit']='V' + y.config['tags']['ytag2'] = 'yyyy' xy=TraceXY('xy trace') - xy.tags['xy tag']= 'I am xy tag' + xy.config['tags']['xy tag']= 'I am xy tag' xy.x = x xy.y = y xy.save('xytrace.dat') - print(xy.getHeader()) + # print(xy.getHeader()) xyn = TraceSetSameX('many ys trace') - xyn.tags['descr'] = 'I am many ys trace' - xy.label='y1' + xyn.config['tags']['descr'] = 'I am many ys trace' + xy.config['label']='y1' xyn.addTrace(xy) - xy.label='y2' + xy.config['label']='y2' xyn.addTrace(xy) - xy.label='y3' + xy.config['label']='y3' xyn.addTrace(xy) xyn.save('xyntrace.dat') - print(xyn.getHeader()) + # print(xyn.getHeader()) |