-
Notifications
You must be signed in to change notification settings - Fork 0
/
stream_backend.py
152 lines (108 loc) · 4.51 KB
/
stream_backend.py
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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
"""
This submodule implements the base class for stream backends.
"""
from __future__ import print_function
import numpy as np
class BackendStream():
"""
Backend for writing to text files
"""
def __init__(self):
self.step = None
self.data_array = None
self.data_list = None
self.draw_limits = None
self.plot_transforms = None
self.plot_options = None
self.interactive = False
self.max_auto_resolution = 1024
self.output_filename = ''
def init_figure(self):
"""
Create the figure
"""
pass
def on_exit(self):
"""
Function to tidy up backend on exit
"""
pass
def set_output_filename(self, filename):
"""
Set the output filename
"""
self.output_filename = filename
def plot_data(self):
"""
Take a data_array and 'plot' to a file
"""
import numpy as np
# Plot limits
xmin, xmax = self.draw_limits['x_axis']
ymin, ymax = self.draw_limits['y_axis']
# Transforms
x_transform = self.plot_transforms['x_transform']
y_transform = self.plot_transforms['y_transform']
hist_transform = self.plot_transforms['hist_transform']
render_transform = self.plot_transforms['render_transform']
limits = self.draw_limits['xy_limits']
self.current_xylimits = limits
plot_type = self.plot_options['plot_type']
if (plot_type == 'time') or (plot_type == 'line_plot'):
self.data_array = self.data_list[0]
elif plot_type == 'hist2d':
xedges, yedges, counts, min_max_data = self.data_list
self.data_array = np.empty((len(xedges)*len(yedges), 3))
k = 0
for i in range(len(xedges)-1):
for j in range(len(xedges)-1):
x_cen = 0.5*(xedges[i] + xedges[i+1])
y_cen = 0.5*(yedges[i] + yedges[i+1])
self.data_array[k, :] = [x_cen, y_cen, counts[i,j]]
k = k + 1
elif plot_type == 'render':
nxy = self.data_list[0].shape
if len(nxy) != 2:
raise ValueError("Didn't do non-3D stuff!")
nx, ny = nxy
half_dx = 0.5 * (xmax - xmin) / float(nx)
half_dy = 0.5 * (ymax - ymin) / float(ny)
x_pos = np.linspace(xmin + half_dx, xmax - half_dx, nx)
y_pos = np.linspace(ymin + half_dy, ymax - half_dy, ny)
self.data_array = np.empty((nx * ny, 3))
k = 0
for i, x in enumerate(x_pos):
for j, y in enumerate(y_pos):
self.data_array[k,:] = [x_pos[i], y_pos[i],
self.data_list[0][i, j]]
k = k + 1
else:
# Single axis plots
data_axis = self.plot_options['data_axis']
if data_axis != 'x':
qx_transform = self.plot_transforms['qx_transform']
else:
qx_transform = None
if data_axis != 'y':
qy_transform = self.plot_transforms['qy_transform']
else:
qy_transform = None
if plot_type.properties['plot_type'] == 'hist1d':
counts, bins, extra_info = self.data_list
if qy_transform is not None:
counts = qy_transform[0](counts)
self.data_array = np.empty((len(bins), 2))
for i, x in enumerate(bins[:-1]):
self.data_array[i, :] = [bins[i], counts[i]]
self.data_array[-1, :] = (bins[-1], 0.0)
elif plot_type.properties['plot_type'] == 'power_spectrum':
x = self.data_list[0]
y = self.data_list[1]
extra_info = self.data_list[2]
if qx_transform is not None:
x = qx_transform[0](x)
if qy_transform is not None:
y = qy_transform[0](y)
self.data_array = np.vstack((x,y)).T
# Write figure to file
self.output_canvas()