#!/usr/bin/env python from __future__ import division import wx import os,re import numpy as np from numpy import apply_along_axis,array,vstack,average,std,sqrt import matplotlib # We want matplotlib to use a wxPython backend matplotlib.use('WXAgg') from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas from matplotlib.figure import Figure from matplotlib.backends.backend_wx import NavigationToolbar2Wx from enthought.traits.api import HasTraits, Any, Instance, Int, Button, Enum from enthought.traits.ui.api import View, Item, ButtonEditor from enthought.traits.ui.wx.editor import Editor #from enthought.traits.ui.wx.basic_editor_factory import BasicEditorFactory from enthought.traits.ui.basic_editor_factory import BasicEditorFactory def splitseq(seq,size): """ Split up seq in pieces of size """ return [seq[i:i+size] for i in range(0, len(seq), size)] def plotdata(x,y=None,yerr=None,numblocks=10,numerrbars=50,colors=('blue','green'),cumavg=True,clear=False, plottitle=None,xaxislabel=None,yaxislabel=None,axes=None): """ Our standard plotting routine. - The first argument is the data. We try to extract x and y data from it if possible. - Plot 50 error bars if we're given errors. - Do block averaging. - Plot the cumulative average. """ if axes is None: raise Error('Need axes instance') txt = '' # Figure out x,y,err if y is None: if len(x.shape) > 1: x,y = x[:,0],x[:,1] else: x,y = array(range(len(x))),x if clear: clf() annotation_location = (min(x) + (max(x) - min(x))*0.1,min(y) + (max(y) - min(y))*0.9) #print annotation_location,max(y) axes.plot(x,y,color=colors[0],zorder=10,alpha=0.5) if cumavg: ya = np.cumsum(y)/np.arange(1,len(y)+1) axes.plot(x,ya,'k-') if yerr is not None: error_step = int(len(x)/numerrbars) if error_step == 0: error_step = len(x) axes.errorbar(x[::error_step],y[::error_step],yerr[::error_step],color=colors[0],zorder=20) blocksize = int(len(x)/numblocks) blocksizes = [len(i) for i in splitseq(x,blocksize)] x = [average(i) for i in splitseq(x,blocksize)][:numblocks] yerr = array([std(i) for i in splitseq(y,blocksize)])[:numblocks] y = array([average(i) for i in splitseq(y,blocksize)])[:numblocks] txt += 'Avg: %.4f'%average(y) txt += ', std err: %.4f, avg err: %.4f'%(std(y)/sqrt(numblocks),average(yerr)) txt += '\nblocks of length %s'%blocksize if len(blocksizes) > numblocks: txt += ' discarding %s points at the end'%sum(blocksizes[numblocks:]) axes.errorbar(x,y,yerr,elinewidth=20,color=colors[1],barsabove=True,zorder=30) if plottitle: axes.set_title(plottitle) if xaxislabel: axes.set_xlabel(xaxislabel) if yaxislabel: axes.set_ylabel(yaxislabel) axes.annotate(txt,annotation_location) class XVG: def __init__(self,fname,startdir='.'): self.fname = os.path.join(startdir,fname) self.attributes = {'comments':''} self.columns = ['time'] self.units = {'time':'ns'} self.process_file() def process_file(self): print "processing",self.fname f = file(self.fname) lines = [] for line in f: if line.startswith('#'): self.attributes['comments'] += line elif line.startswith('@'): self.add_info(line) else: lines.append([float(i) for i in line.split()]) f.close() if lines and (len(lines[-1]) != len(lines[-2])): print "Deleting bad last line :",lines[-1] print "Expected something like:",lines[-2] lines = lines[:-1] self.data = array(lines) if len(self.data) > 0: self.data[::,0] = self.data[::,0]/1000. # convert to nanoseconds def add_info(self,line): line = line[1:].strip() info_type = line.split()[0] rest = line[line.index(info_type)+len(info_type):].strip() if info_type in 'title TYPE'.split(): self.attributes[info_type] = rest.replace('"','').strip() elif info_type in 'xaxis yaxis'.split(): self.attributes[info_type] = rest.replace('label','').replace('"','').strip() if info_type == 'xaxis' and self.attributes[info_type] == 'Time (ps)': self.attributes[info_type] = 'Time (ns)' if info_type == 'yaxis': self.defaultunit = rest.split('(')[-1].replace(')','').replace('"','').strip() elif re.match('s\d+$',info_type): #@ s1 legend "Coulomb (SR)" rest = rest.split('"')[-2] #Coulomb (SR) if '(' in rest: label = rest.split('(')[0].strip() unit = rest.split('(')[-1].replace(')','').strip() self.columns.append(label) self.units[label] = unit else: self.columns.append(rest) self.units[rest] = self.defaultunit#'kJ mol\S-1\N' def get(self,item): result = self.data[:,self.columns.index(item)] return result def getxy(self,item): return self.data[:,0],self.get(item) def plot(self,axes,item,color='',truncat=None, numblocks=10,deavg=False): x,y=self.getxy(item) if truncat: x = x[:truncat] y = y[:truncat] if deavg: print "Subtracting averages y = %s"%(average(y),) y = y - average(y) plotdata(x,y,numblocks=numblocks,xaxislabel=self.attributes['xaxis'],yaxislabel='%s (%s)'%(item,self.units[item]),plottitle=item+' ('+os.path.split(self.fname)[-1]+')',axes=axes) def plotitem(self,idx,axes): self.plot(axes,self.columns[idx]) def plotall(self,figure,ncol=4,tstart=None,skipTime=False): cols = self.columns[:] if skipTime: cols.remove('time') self.plotseveral(figure,cols,ncol) def plotseveral(self,figure,groups,ncol): nrow = np.ceil(len(groups)/ncol) for (i,item) in enumerate(groups): print "Plotting",i,"out of",len(groups),"in subplot",(nrow,ncol,i+1) axes = figure.add_subplot(nrow,ncol,i+1) self.plot(axes,item) # MPL classes taken from Traits tutorial http://code.enthought.com/projects/traits/docs/html/index.html class _MPLFigureEditor(Editor): scrollable = True def init(self, parent): self.control = self._create_canvas(parent) self.set_tooltip() def update_editor(self): pass def _create_canvas(self, parent): """ Create the MPL canvas. """ # The panel lets us add additional controls. panel = wx.Panel(parent, -1, style=wx.CLIP_CHILDREN) sizer = wx.BoxSizer(wx.VERTICAL) panel.SetSizer(sizer) # matplotlib commands to create a canvas mpl_control = FigureCanvas(panel, -1, self.value) sizer.Add(mpl_control, 1, wx.LEFT | wx.TOP | wx.GROW) toolbar = NavigationToolbar2Wx(mpl_control) sizer.Add(toolbar, 0, wx.EXPAND) self.value.canvas.SetMinSize((10,10)) return panel class MPLFigureEditor(BasicEditorFactory): klass = _MPLFigureEditor if __name__ == "__main__": # Create a window to demo the editor from numpy import sin, cos, linspace, pi usage='''Just a simple XVG viewer. For each column in the XVG file, we plot - the raw data (light blue) - a running average (black) - block averages (green) We also print out - the average over the simulation - standard error over the blocks - average of the standard deviations within each block Block length is reported in number of frames, and its meaning will therefore depend on the settings in your .mdp file as well as the arguments you used when making the .xvg file. We attempt to get units out of the XVG file, and we convert times to nanoseconds before plotting. This program requires - Python 2.5 or 2.6 - wxPython 2.8.7.1+ (http://www.wxpython.org/) - numpy 1.3.0+ (http://numpy.scipy.org/) - matplotlib 0.99.0+ (http://matplotlib.sourceforge.net/) - traits 3.2.0+ (http://code.enthought.com/projects/traits/) All of which will be installed for you if you use the Enthought Python Distribution (http://enthought.com/). I mostly wrote this as a quick way to figure out how to use the Traits GUI, but feel free to email me (m g lerner atsign gmail dot com) if you have any questions. ''' from optparse import OptionParser parser = OptionParser(usage=usage) parser.add_option("-f", "--file", help="The XVG file [default %default]", default='energy.xvg') (options, args) = parser.parse_args() class Test(HasTraits): figure = Instance(Figure, ()) value = 0 plot_next = Button() plot_prev = Button() plot_all = Button() xvg = XVG(options.file) items = Enum(*xvg.columns) def _plot_next_fired(self): self.value +=1 if self.value >= len(self.xvg.columns): self.value = 0 self.items = self.xvg.columns[self.value] # auto-calls _items_changed def _plot_prev_fired(self): self.value -=1 if self.value < 0: self.value = len(self.xvg.columns) - 1 self.items = self.xvg.columns[self.value] # auto-calls _items_changed def _plot_all_fired(self): self.figure.clear() self.xvg.plotall(self.figure,skipTime=True) self.figure.canvas.draw() def plot_curr_item(self): self.figure.clear() axes = self.figure.add_subplot(111) self.xvg.plotitem(self.value,axes) self.figure.canvas.draw() def _items_changed(self,old,new): self.value = self.xvg.columns.index(new) self.plot_curr_item() view = View(Item('plot_prev', show_label=False), Item('plot_next', show_label=False), Item('plot_all', show_label=False), Item('items', show_label=False), Item('figure', editor=MPLFigureEditor(), show_label=False), width=800, height=600, resizable=True, ) def __init__(self): super(Test, self).__init__() # Start off with a cute flower axes = self.figure.add_subplot(111) t = linspace(0, 2*pi, 200) axes.plot(sin(t)*(1+0.5*cos(11*t)), cos(t)*(1+0.5*cos(11*t))) axes.annotate('Please click a button or \nuse the dropdown menu',(-0.37,0.)) Test().configure_traits()