aboutsummaryrefslogtreecommitdiff
path: root/qolab/data/trace.py
blob: dccc2f73b4aeb9b0e55c2849fb1724dd754ccd54 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
from qolab.file_utils import save_table_with_header

class Trace:
    def __init__(self, label):
        self.label = label
        self.values = None
        self.unit = None
        self.tags = {}

    def plot(self):
        import matplotlib.pyplot as plt
        plt.plot(self.values, label=self.descr)
        plt.legend()
    
    def header(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}')
        prefixed_header = [prefix+l for l in header]
        return prefixed_header

    def write(self, fname, item_format='e'):
        save_table_with_header(fname, self.values, self.header(), item_format=item_format)


class TraceXY:
    def __init__(self, label):
        self.label = label
        self.tags = {}
        self.x = None
        self.y = None

    def plot(self):
        import matplotlib.pyplot as plt
        plt.plot(self.x.values, self.y.values, label=self.label)
        plt.legend()
        plt.xlabel(self.x.label)

    def header(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.header(prefix='x_info: '))
        header.extend(self.y.header(prefix='y_info: '))
        prefixed_header = [prefix+l for l in header]
        return prefixed_header
        
class TraceSetSameX:
    def __init__(self, label):
        self.label = label
        self.tags = {}
        self.x = None
        self.traces={}

    def addTrace(self, tr):
        if len(self.traces) == 0:
            self.x = tr.x
        trY = tr.y
        self.traces[tr.label]=trY
        
    def plot(self):
        import matplotlib.pyplot as plt
        for k, tr in self.traces.items():
            plt.plot(self.x.values, tr.values, label=tr.label)
        plt.xlabel(self.x.label)
        plt.legend()

    def items(self):
        return (self.traces.items())

    def keys(self):
        return (self.traces.keys())

    def getTrace(self, label):
        tr = TraceXY(label)
        tr.x = self.x
        tr.y = self.traces[label]
        return (tr)
    
    def header(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.header(prefix='x_info: '))
        for k, v in self.traces.items():
            header.extend(v.header(prefix=f'{k}_info: '))
        prefixed_header = [prefix+l for l in header]
        return prefixed_header
        

if __name__ == '__main__': 
    import numpy as np
    print("Testing trace")
    x=Trace('x trace')
    x.values = np.random.normal(2,2,(4,2))
    x.unit='s'
    x.tags['x_tag'] = 'xxxx'
    x.tags['x_tag2'] = 'xxxx'
    x.write('xtrace.dat')
    y=Trace('y trace')
    y.values = np.random.normal(2,2,(4,2))
    y.unit='V'
    y.tags['y_tag'] = 'yyyy'
    xy=TraceXY('xy trace')
    xy.tags['xy tag']= 'I am xy tag'
    xy.x = x
    xy.y = y
    print(xy.header())
    xyn = TraceSetSameX('many ys trace')
    xyn.tags['descr'] = 'I am many ys trace'
    xy.label='y1'
    xyn.addTrace(xy)
    xy.label='y2'
    xyn.addTrace(xy)
    xy.label='y3'
    xyn.addTrace(xy)
    print(xyn.header())