-
Notifications
You must be signed in to change notification settings - Fork 3
/
plot_generator.py
117 lines (96 loc) · 3.63 KB
/
plot_generator.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
from traits.api import Str, HasTraits, Button, Enum, Bool
from traitsui.api import UItem, HGroup, VGroup, Group, View, spring
from processing import rescale, stack_datasets
from mpl_plot import MplPlot
from chaco_plot import StackedPlot, Surface2DPlot
from ui_helpers import get_save_as_filename, open_file_with_default_handler
class PlotGenerator(HasTraits):
plot_types = [
"Stacked",
"2D surface",
"3D plot (slow)",
]
plot_type = Enum(plot_types)
reset_button = Button("Reset view")
save_button = Button("Save plot...")
copy_button = Button("Copy to clipboard")
scale = Enum('linear', 'log', 'sqrt')
replot = Bool(False)
status = Str
def _on_redraw(self, drawing):
self.status = 'Rendering plot...' if drawing else 'Done'
def _set_plots(self, *args):
plot_list = args
self.plots = {}
for i, plot in enumerate(plot_list):
self.plots[self.plot_types[i]] = plot
# Traits needs an attribute set on this class to
# be able to find traits in the plot objects.
setattr(self, self._plot_attr_name(plot), plot)
def _plot_attr_name(self, plot):
return '_plot_' + str(id(plot))
def __init__(self, *args, **kws):
self._set_plots(
StackedPlot(),
Surface2DPlot(),
MplPlot(callback_obj=self),
)
super(PlotGenerator, self).__init__(*args, **kws)
self.datasets = kws['datasets']
self.stack = stack_datasets(self.datasets)
self.cached_data = {}
self._plot_type_changed()
def __del__(self):
del self.plots
def show(self):
menu_group = HGroup(
UItem('plot_type', style='custom'),
spring,
UItem('scale'),
UItem('reset_button'),
spring,
UItem('save_button'),
UItem('copy_button'),
spring,
)
plot_containers = []
for plot_type in PlotGenerator.plot_types:
plot_group = self.plots[plot_type].traits_group(
self._plot_attr_name(self.plots[plot_type]),
visible_when="plot_type == '%s'" % plot_type
)
plot_containers.append(plot_group)
view = View(
VGroup(
menu_group,
Group(*plot_containers),
),
title='Plot generator', resizable=True, width=900, height=600,
statusbar='status',
)
self.edit_traits(view=view)
def _scale_changed(self):
self.replot = True
self._plot_type_changed()
self.plots[self.plot_type].reset_view()
def _reset_button_changed(self):
self.plots[self.plot_type].reset_view()
def _save_button_changed(self):
filename = get_save_as_filename()
if filename is None:
return
self.plots[self.plot_type].save_as(filename)
open_file_with_default_handler(filename)
def _copy_button_changed(self):
self.plots[self.plot_type].copy_to_clipboard()
def _plot_type_changed(self):
self.replot = True
self._generate_plot()
def _generate_plot(self):
if (self.plot_type not in self.cached_data) or self.replot:
x, y, z = self.plots[self.plot_type].prepare_data(self.stack)
z = rescale(z, method=self.scale)
self.cached_data[self.plot_type] = x, y, z
self.replot = False
x, y, z = self.cached_data[self.plot_type]
return self.plots[self.plot_type].plot(x, y, z, scale=self.scale)