$search
00001 #!/usr/bin/env python 00002 # 00003 # Based on public domain code from Eli Bendersky (eliben@gmail.com) 00004 # http://eli.thegreenplace.net/2008/08/01/matplotlib-with-wxpython-guis/ 00005 # 00006 # Additions are covered by the license below: 00007 # 00008 # Software License Agreement (BSD License) 00009 # 00010 # Copyright (c) 2009, Willow Garage, Inc. 00011 # All rights reserved. 00012 # 00013 # Redistribution and use in source and binary forms, with or without 00014 # modification, are permitted provided that the following conditions 00015 # are met: 00016 # 00017 # * Redistributions of source code must retain the above copyright 00018 # notice, this list of conditions and the following disclaimer. 00019 # * Redistributions in binary form must reproduce the above 00020 # copyright notice, this list of conditions and the following 00021 # disclaimer in the documentation and/or other materials provided 00022 # with the distribution. 00023 # * Neither the name of Willow Garage, Inc. nor the names of its 00024 # contributors may be used to endorse or promote products derived 00025 # from this software without specific prior written permission. 00026 # 00027 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 00028 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 00029 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS 00030 # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE 00031 # COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, 00032 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 00033 # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 00034 # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 00035 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 00036 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN 00037 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 00038 # POSSIBILITY OF SUCH DAMAGE. 00039 # 00040 # Revision $Id$ 00041 00042 PKG = 'rxtools' 00043 import collections 00044 import os 00045 import string 00046 import sys 00047 import wx 00048 00049 import roslib.names 00050 import roslib.packages 00051 00052 import rospy 00053 00054 import rxtools.vizutil 00055 rxtools.vizutil.check_matplotlib_deps() 00056 00057 ## MPL imports 00058 00059 import matplotlib 00060 matplotlib.use('WXAgg') 00061 from matplotlib.figure import Figure 00062 from matplotlib.backends.backend_wxagg import \ 00063 FigureCanvasWxAgg as FigCanvas, \ 00064 NavigationToolbar2WxAgg as NavigationToolbar 00065 00066 import numpy as np 00067 import pylab 00068 00069 from rxtools.rosplot import ROSData, is_ros_pause, is_ros_stop, toggle_ros_pause, set_ros_stop 00070 import mpl_toolkits.mplot3d.art3d as art3d 00071 00072 ## Plotting colors. We artificially constrain the number of allowed plots by the number of 00073 ## COLORS we specify. rxplot has not been stress-tested yet 00074 COLORS = (1, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 1), (0, 1, 1), (0.5, 0.24, 0), (0, 0.5, 0.24), (1, 0.5, 0), 00075 00076 # TODO: rewrite so that pause button toggles activation of other buttons 00077 ## Toolbar controls for rxplot 00078 class RxPlotToolbar(NavigationToolbar): 00079 ON_PAUSE = wx.NewId() 00080 ON_STOP = wx.NewId() 00081 def __init__(self, canvas): 00082 NavigationToolbar.__init__(self, canvas) 00083 00084 self.play_bm = wx.Bitmap(roslib.packages.get_pkg_dir(PKG) + '/icons/play-16.png') 00085 self.pause_bm = wx.Bitmap(roslib.packages.get_pkg_dir(PKG) + '/icons/pause-16.png') 00086 self.AddSimpleTool(self.ON_PAUSE, self.pause_bm, 'Pause', 'Activate pause') 00087 wx.EVT_TOOL(self, self.ON_PAUSE, self._on_pause) 00088 00089 self.AddSimpleTool(self.ON_STOP, wx.Bitmap(roslib.packages.get_pkg_dir(PKG) + '/icons/stop-16.png'), 'Stop', 'Activate stop') 00090 00091 wx.EVT_TOOL(self, self.ON_STOP, self._on_stop) 00092 self._set_tb_enable() 00093 00094 ## Most toolbar buttons remain disabled until rxplot enters paused or stopped state. 00095 def _set_tb_enable(self): 00096 # parent toolbar doesn't expose home, save, or subplot 00097 # buttons. However, these don't behave as badly under pause so 00098 # it's not as important to disable those. 00099 tb_enabled = is_ros_pause() or is_ros_stop() 00100 for b in [self._NTB2_BACK, self._NTB2_FORWARD, self._NTB2_PAN, self._NTB2_ZOOM]: 00101 self.EnableTool(b, tb_enabled) 00102 00103 def _on_pause(self, evt): 00104 toggle_ros_pause() 00105 self._set_tb_enable() 00106 if is_ros_pause(): 00107 self.SetToolNormalBitmap(self.ON_PAUSE, self.play_bm) 00108 else: 00109 self.SetToolNormalBitmap(self.ON_PAUSE, self.pause_bm) 00110 # cancel state of pan/zoom 00111 self.ToggleTool(self._NTB2_PAN, False) 00112 self.ToggleTool(self._NTB2_ZOOM, False) 00113 self._active = 'ZOOM' 00114 self.zoom(tuple()) 00115 00116 def _on_stop(self, evt): 00117 set_ros_stop() 00118 self._set_tb_enable() 00119 # kill pause/stop buttons as well 00120 self.EnableTool(self.ON_PAUSE, False) 00121 self.EnableTool(self.ON_STOP, False) 00122 00123 ## WX Frame containing plots 00124 class RxPlotFrame(wx.Frame): 00125 def __init__(self, topics, options): 00126 self.title = options.title 00127 self.legend = options.legend 00128 self.mode = options.mode 00129 if self.mode == '2d': 00130 self.draw_plot = self.draw_plot_2d 00131 self.init_plot = self.init_plot_2d 00132 elif self.mode in ['3d', 'scatter']: 00133 self.draw_plot = self.draw_plot_3d 00134 self.init_plot = self.init_plot_3d 00135 else: 00136 raise ValueError("Mode must be one of [2d, 3d, scatter], not [%s]"%(options.mode)) 00137 00138 wx.Frame.__init__(self, None, -1, self.title) 00139 00140 # There are two different indicies for our list data types based on the fact that topics is 00141 # a list[list(str)]. 00142 # 00143 # axes_index: first index into topics 00144 # 00145 # plot_index: index based on flattened topic list, i.e. absolute topic number 00146 # 00147 # e.g. for self.topics = [ ['/p0/f0'], ['/p1/f1', '/p2/f2']]: 00148 # * axes_index 0 indexes ['/p0/f0'], 00149 # * axes_index 1 indexes ['/p1/f1', '/p2/f2'], 00150 # * plot_index 0 indexes '/p0/f0' 00151 # * plot_index 1 indexes '/p1/f1' 00152 # * plot_index 2 indexes '/p2/f2' 00153 00154 # self.axes and self.topics uses axes_index 00155 self.axes = [] 00156 self.topics = topics 00157 self.period = options.period 00158 self.marker = options.marker 00159 00160 # initialize data storage and sources 00161 # datagen, datax, and datay use plot_index. 00162 self.datagen = [] 00163 self.datax = [] 00164 self.datay = [] 00165 self.buffer_size = options.buffer 00166 self.redraw_period = int(1. / options.refresh_rate * 1000) # convert to ms 00167 00168 self.start_time = rospy.get_time() 00169 for topic_list in topics: 00170 for t in topic_list: 00171 dg = ROSData(t, self.start_time) 00172 self.datagen.append(dg) 00173 datax, datay = dg.next() 00174 self.datax.append(collections.deque(datax)) 00175 self.datay.append(collections.deque(datay)) 00176 00177 self.create_menu() 00178 self.statusbar = self.CreateStatusBar() 00179 self.create_main_canvas() 00180 00181 self.redraw_timer = wx.Timer(self) 00182 self.Bind(wx.EVT_TIMER, self.on_redraw_timer, self.redraw_timer) 00183 00184 #self.redraw_timer.Start(100) 00185 self.redraw_timer.Start(self.redraw_period) 00186 if options.start_paused: 00187 self.toolbar._on_pause(None) 00188 00189 def create_menu(self): 00190 self.menubar = wx.MenuBar() 00191 00192 menu_file = wx.Menu() 00193 m_expt = menu_file.Append(-1, "&Save plot\tCtrl-S", "Save plot to file") 00194 self.Bind(wx.EVT_MENU, self.on_save_plot, m_expt) 00195 menu_file.AppendSeparator() 00196 m_exit = menu_file.Append(-1, "E&xit\tCtrl-X", "Exit") 00197 self.Bind(wx.EVT_MENU, self.on_exit, m_exit) 00198 00199 self.menubar.Append(menu_file, "&File") 00200 self.SetMenuBar(self.menubar) 00201 00202 def onpick(self, event): 00203 thisline = event.artist 00204 xdata, ydata = thisline.get_data() 00205 ind = event.ind 00206 self.statusbar.SetStatusText("%s, %s" % (xdata[ind][len(xdata[ind])/2], ydata[ind][len(ydata[ind])/2])) 00207 00208 def create_main_canvas(self): 00209 # Create main figure 00210 self.dpi = 100 00211 params = matplotlib.figure.SubplotParams(left=0.125, bottom=0.12, right=0.99, top=0.99, wspace=0.001, hspace=0.1) 00212 self.fig = Figure((3.0, 3.0), dpi=self.dpi, subplotpars=params) 00213 00214 self.canvas = FigCanvas(self, -1, self.fig) 00215 self.canvas.mpl_connect('pick_event', self.onpick) 00216 00217 self.init_plot() 00218 00219 self.sizer = wx.BoxSizer(wx.VERTICAL) 00220 self.sizer.Add(self.canvas, 1, flag=wx.LEFT | wx.TOP | wx.GROW) 00221 self.SetSizer(self.sizer) 00222 00223 self.add_toolbar() 00224 00225 self.SetSize(wx.Size(500, 700)) 00226 self.Layout() 00227 00228 def init_plot_2d(self): 00229 self.plot_data = [] 00230 flat_topics = [] 00231 00232 axes_index = 0 00233 plot_index = 0 00234 00235 for topic_list in self.topics: 00236 axes = self.fig.add_subplot(string.atoi('%d1%d'%(len(self.topics), axes_index+1))) 00237 axes.set_axis_bgcolor('white') 00238 #axes.set_title(', '.join(topic_list), size=8) 00239 pylab.setp(axes.get_xticklabels(), fontsize=6) 00240 pylab.setp(axes.get_yticklabels(), fontsize=8) 00241 # equiv to self.axes[axes_index] = axes 00242 self.axes.append(axes) 00243 axes_index += 1 00244 00245 for topic in topic_list: 00246 flat_topics.append(topic) 00247 plot_data = \ 00248 axes.plot(self.datax[plot_index], 00249 self.datay[plot_index], 00250 marker=self.marker, 00251 linewidth=1, 00252 picker=5, 00253 color=COLORS[plot_index], 00254 )[0] 00255 # equiv to self.plot_data[plot_index] = plot_data 00256 self.plot_data.append(plot_data) 00257 plot_index += 1 00258 00259 # Set grid visibility properties 00260 for ax in self.axes: 00261 ax.grid(True, color='gray') 00262 pylab.setp(ax.get_xticklabels(), visible=True) 00263 00264 # Make a legend 00265 fp = matplotlib.font_manager.FontProperties(size=8) 00266 legends = self.legend.split(",") if self.legend else [] 00267 if legends: 00268 if len(legends) != len(flat_topics): 00269 raise "Number of legend strings does not match the number of topics" 00270 self.fig.legend(self.plot_data, legends, 'lower right', prop=fp) 00271 elif len(flat_topics) > 1: 00272 self.fig.legend(self.plot_data, flat_topics, 'lower right', prop=fp) 00273 00274 def init_plot_3d(self): 00275 # init plot_data (only used for 3d mode right now until scatter is ported) 00276 self.plot_data = [] 00277 # choose a color (for scatter mode only) 00278 import random 00279 self.color = COLORS[random.randint(0, len(COLORS))] 00280 00281 # initialize 3d axes from mplot3d 00282 import mpl_toolkits.mplot3d.axes3d 00283 self.ax = mpl_toolkits.mplot3d.axes3d.Axes3D(self.fig) 00284 00285 flat_topics = [] 00286 for topic_list in self.topics: 00287 for topic in topic_list: 00288 flat_topics.append(topic) 00289 00290 ntopics = len(flat_topics) 00291 if ntopics <= 3: 00292 self.ax.set_xlabel(flat_topics[0]) 00293 self.ax.set_ylabel(flat_topics[1]) 00294 if ntopics == 2: 00295 self.ax.set_zlabel("time") 00296 elif ntopics == 3: 00297 self.ax.set_zlabel(flat_topics[2]) 00298 else: 00299 raise Exception("Expected 2 or 3 topics, but got %d" % (ntopics)) 00300 00301 # attach mouse handler (must be called after FigCanvas) 00302 self.ax.mouse_init() 00303 00304 def draw_plot_2d(self, relimit=False): 00305 if not self.plot_data: 00306 return 00307 00308 # Set axis bounds 00309 if relimit and self.datax[0]: 00310 axes_index = 0 00311 plot_index = 0 00312 # axes are indexed by topic_list, plots are indexed by topic number 00313 for topic_list in self.topics: 00314 axes = self.axes[axes_index] 00315 axes_index += 1 00316 ymin = ymax = None 00317 for t in topic_list: 00318 datax = self.datax[plot_index] 00319 datay = self.datay[plot_index] 00320 plot_index += 1 00321 00322 xmax = datax[-1] 00323 xmin = xmax - self.period 00324 if ymin is None: 00325 ymin = min(datay) 00326 ymax = max(datay) 00327 else: 00328 ymin = min(min(datay), ymin) 00329 ymax = max(max(datay), ymax) 00330 00331 # pad the min/max 00332 delta = ymax - ymin 00333 ymin -= .05*delta 00334 ymax += .05*delta 00335 00336 axes.set_xbound(lower=xmin, upper=xmax) 00337 axes.set_ybound(lower=ymin, upper=ymax) 00338 00339 # Set plot data on current axes 00340 for plot_index in xrange(0, len(self.plot_data)): 00341 datax = self.datax[plot_index] 00342 datay = self.datay[plot_index] 00343 00344 plot_data = self.plot_data[plot_index] 00345 plot_data.set_data(np.array(datax), np.array(datay)) 00346 00347 self.canvas.draw() 00348 00349 def draw_plot_3d(self, relimit=False): 00350 # Don't have to relimit 3d plots 00351 if self.datax[0] and self.datay[0] and self.datay[1]: 00352 00353 ndata = len(self.datay) 00354 if ndata >=2: 00355 datax = np.array(self.datay[0]) 00356 datay = np.array(self.datay[1]) 00357 00358 # Plot time from topic 1 on z if no third topic is provided 00359 if ndata == 2: 00360 dataz = np.array(self.datax[0]) 00361 elif ndata == 3 and self.datay[2]: 00362 dataz = np.array(self.datay[2]) 00363 else: 00364 print "unexpected input data dimensions for 3d plotting (%d)" % ndata 00365 00366 if self.mode == '3d': 00367 if not self.plot_data: 00368 self.plot_data = self.ax.plot(datax, datay, dataz) 00369 00370 self.plot_data[0].set_data(datax, datay) 00371 art3d.line_2d_to_3d(self.plot_data[0], zs=dataz, zdir='z') 00372 self.ax.auto_scale_xyz(datax, datay, dataz, True) 00373 else: 00374 # 'blue' is arbitrary, just have to keep mpl from choosing different each time. 00375 # 00376 # kwc: I haven't figured out how to performance tune 00377 # scatter plots the same way I tuned the 3d plot. The 00378 # MPL APIs are a bit haphazard when it comes to 3d 00379 # plotting. 00380 self.plot_data = self.ax.scatter(datax, datay, dataz, color=self.color) 00381 00382 self.canvas.draw() 00383 00384 def on_save_plot(self, event): 00385 file_choices = "PNG (*.png)|*.png" 00386 00387 dlg = wx.FileDialog( 00388 self, 00389 message="Save plot as...", 00390 defaultDir=os.getcwd(), 00391 defaultFile="plot.png", 00392 wildcard=file_choices, 00393 style=wx.SAVE) 00394 00395 if dlg.ShowModal() == wx.ID_OK: 00396 path = dlg.GetPath() 00397 self.canvas.print_figure(path, dpi=self.dpi) 00398 self.flash_status_message("Saved to %s" % path) 00399 00400 def on_redraw_timer(self, event): 00401 # if paused do not add data, but still redraw the plot 00402 # (to respond to scale modifications, grid change, etc.) 00403 # 00404 if is_ros_stop(): 00405 for dg in self.datagen: 00406 dg.close() 00407 self.datagen = [] 00408 00409 if not is_ros_pause(): 00410 for plot_index in xrange(0, len(self.datagen)): 00411 try: 00412 datax, datay = self.datagen[plot_index].next() 00413 except Exception, e: 00414 print >> sys.stderr, str(e) 00415 wx.GetApp().Exit() 00416 return 00417 00418 plot_datax, plot_datay = self.datax[plot_index], self.datay[plot_index] 00419 plot_datax.extend(datax) 00420 plot_datay.extend(datay) 00421 00422 if len(plot_datax) > 0 and self.buffer_size > 0: 00423 xcutoff = plot_datax[-1] - self.buffer_size 00424 while len(plot_datax) > 0 and plot_datax[0] < xcutoff: 00425 plot_datax.popleft() 00426 plot_datay.popleft() 00427 00428 self.draw_plot(relimit=True) 00429 00430 def on_exit(self, event): 00431 self.Destroy() 00432 00433 def flash_status_message(self, msg, flash_len_ms=1500): 00434 self.statusbar.SetStatusText(msg) 00435 self.timeroff = wx.Timer(self) 00436 self.Bind( 00437 wx.EVT_TIMER, 00438 self.on_flash_status_off, 00439 self.timeroff) 00440 self.timeroff.Start(flash_len_ms, oneShot=True) 00441 00442 def on_flash_status_off(self, event): 00443 self.statusbar.SetStatusText('') 00444 00445 def add_toolbar(self): 00446 self.toolbar = RxPlotToolbar(self.canvas) 00447 self.toolbar.Realize() 00448 self.SetToolBar(self.toolbar) 00449 if 0: 00450 if wx.Platform == '__WXMAC__': 00451 # Mac platform (OSX 10.3, MacPython) does not seem to cope with 00452 # having a toolbar in a sizer. This work-around gets the buttons 00453 # back, but at the expense of having the toolbar at the top 00454 self.SetToolBar(self.toolbar) 00455 else: 00456 # On Windows platform, default window size is incorrect, so set 00457 # toolbar width to figure width. 00458 tw, th = self.toolbar.GetSizeTuple() 00459 fw, fh = self.canvas.GetSizeTuple() 00460 # By adding toolbar in sizer, we are able to put it at the bottom 00461 # of the frame - so appearance is closer to GTK version. 00462 # As noted above, doesn't work for Mac. 00463 self.toolbar.SetSize(wx.Size(fw, th)) 00464 self.sizer.Add(self.toolbar, 0, wx.LEFT | wx.EXPAND) 00465 # update the axes menu on the toolbar 00466 self.toolbar.update() 00467 00468 00469 00470 def rxplot_app(topic_list, options): 00471 rospy.init_node('rxplot', anonymous=True) 00472 00473 try: 00474 app = wx.PySimpleApp() 00475 app.frame = rxtools.rxplot.RxPlotFrame(topic_list, options) 00476 app.frame.Show() 00477 app.MainLoop() 00478 except Exception, e: 00479 rospy.logerr(e) 00480 print >> sys.stderr, str(e) 00481 wx.GetApp().Exit() 00482 00483 rospy.signal_shutdown('GUI shutdown') 00484