barchart.py revision 1881
1# Copyright (c) 2005 The Regents of The University of Michigan 2# All rights reserved. 3# 4# Redistribution and use in source and binary forms, with or without 5# modification, are permitted provided that the following conditions are 6# met: redistributions of source code must retain the above copyright 7# notice, this list of conditions and the following disclaimer; 8# redistributions in binary form must reproduce the above copyright 9# notice, this list of conditions and the following disclaimer in the 10# documentation and/or other materials provided with the distribution; 11# neither the name of the copyright holders nor the names of its 12# contributors may be used to endorse or promote products derived from 13# this software without specific prior written permission. 14# 15# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 16# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 17# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 18# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 19# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 20# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 21# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 22# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 23# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26# 27# Authors: Nathan Binkert 28# Lisa Hsu 29 30import matplotlib, pylab 31from matplotlib.numerix import array, arange, reshape, shape, transpose, zeros 32from matplotlib.numerix import Float 33 34matplotlib.interactive(False) 35 36class BarChart(object): 37 def __init__(self, **kwargs): 38 self.init(**kwargs) 39 40 def init(self, **kwargs): 41 self.colormap = 'jet' 42 self.inputdata = None 43 self.chartdata = None 44 self.xlabel = None 45 self.ylabel = None 46 self.legend = None 47 self.xticks = None 48 self.yticks = None 49 self.title = None 50 51 for key,value in kwargs.iteritems(): 52 self.__setattr__(key, value) 53 54 def gen_colors(self, count): 55 cmap = matplotlib.cm.get_cmap(self.colormap) 56 if count == 1: 57 return cmap([ 0.5 ]) 58 else: 59 return cmap(arange(count) / float(count - 1)) 60 61 # The input data format does not match the data format that the 62 # graph function takes because it is intuitive. The conversion 63 # from input data format to chart data format depends on the 64 # dimensionality of the input data. Check here for the 65 # dimensionality and correctness of the input data 66 def set_data(self, data): 67 if data is None: 68 self.inputdata = None 69 self.chartdata = None 70 return 71 72 data = array(data) 73 dim = len(shape(data)) 74 if dim not in (1, 2, 3): 75 raise AttributeError, "Input data must be a 1, 2, or 3d matrix" 76 self.inputdata = data 77 78 # If the input data is a 1d matrix, then it describes a 79 # standard bar chart. 80 if dim == 1: 81 self.chartdata = array([[data]]) 82 83 # If the input data is a 2d matrix, then it describes a bar 84 # chart with groups. The matrix being an array of groups of 85 # bars. 86 if dim == 2: 87 self.chartdata = transpose([data], axes=(2,0,1)) 88 89 # If the input data is a 3d matrix, then it describes an array 90 # of groups of bars with each bar being an array of stacked 91 # values. 92 if dim == 3: 93 self.chartdata = transpose(data, axes=(1,2,0)) 94 95 def get_data(self): 96 return self.inputdata 97 98 data = property(get_data, set_data) 99 100 # Graph the chart data. 101 # Input is a 3d matrix that describes a plot that has multiple 102 # groups, multiple bars in each group, and multiple values stacked 103 # in each bar. The underlying bar() function expects a sequence of 104 # bars in the same stack location and same group location, so the 105 # organization of the matrix is that the inner most sequence 106 # represents one of these bar groups, then those are grouped 107 # together to make one full stack of bars in each group, and then 108 # the outer most layer describes the groups. Here is an example 109 # data set and how it gets plotted as a result. 110 # 111 # e.g. data = [[[10,11,12], [13,14,15], [16,17,18], [19,20,21]], 112 # [[22,23,24], [25,26,27], [28,29,30], [31,32,33]]] 113 # 114 # will plot like this: 115 # 116 # 19 31 20 32 21 33 117 # 16 28 17 29 18 30 118 # 13 25 14 26 15 27 119 # 10 22 11 23 12 24 120 # 121 # Because this arrangement is rather conterintuitive, the rearrange 122 # function takes various matricies and arranges them to fit this 123 # profile. 124 # 125 # This code deals with one of the dimensions in the matrix being 126 # one wide. 127 # 128 def graph(self): 129 if self.chartdata is None: 130 raise AttributeError, "Data not set for bar chart!" 131 132 self.figure = pylab.figure() 133 self.axes = self.figure.add_subplot(111) 134 135 dim = len(shape(self.inputdata)) 136 cshape = shape(self.chartdata) 137 if dim == 1: 138 colors = self.gen_colors(cshape[2]) 139 colors = [ [ colors ] * cshape[1] ] * cshape[0] 140 141 if dim == 2: 142 colors = self.gen_colors(cshape[0]) 143 colors = [ [ [ c ] * cshape[2] ] * cshape[1] for c in colors ] 144 145 if dim == 3: 146 colors = self.gen_colors(cshape[1]) 147 colors = [ [ [ c ] * cshape[2] for c in colors ] ] * cshape[0] 148 149 colors = array(colors) 150 151 bars_in_group = len(self.chartdata) 152 if bars_in_group < 5: 153 width = 1.0 / ( bars_in_group + 1) 154 center = width / 2 155 else: 156 width = .8 / bars_in_group 157 center = .1 158 159 bars = [] 160 for i,stackdata in enumerate(self.chartdata): 161 bottom = array([0] * len(stackdata[0])) 162 stack = [] 163 for j,bardata in enumerate(stackdata): 164 bardata = array(bardata) 165 ind = arange(len(bardata)) + i * width + center 166 bar = self.axes.bar(ind, bardata, width, bottom=bottom, 167 color=colors[i][j]) 168 stack.append(bar) 169 bottom += bardata 170 bars.append(stack) 171 172 if self.xlabel is not None: 173 self.axes.set_xlabel(self.xlabel) 174 175 if self.ylabel is not None: 176 self.axes.set_ylabel(self.ylabel) 177 178 if self.yticks is not None: 179 ymin, ymax = self.axes.get_ylim() 180 nticks = float(len(self.yticks)) 181 ticks = arange(nticks) / (nticks - 1) * (ymax - ymin) + ymin 182 self.axes.set_yticks(ticks) 183 self.axes.set_yticklabels(self.yticks) 184 185 if self.xticks is not None: 186 self.axes.set_xticks(arange(cshape[2]) + .5) 187 self.axes.set_xticklabels(self.xticks) 188 189 if self.legend is not None: 190 if dim == 1: 191 lbars = bars[0][0] 192 if dim == 2: 193 lbars = [ bars[i][0][0] for i in xrange(len(bars))] 194 if dim == 3: 195 number = len(bars[0]) 196 lbars = [ bars[0][number - j - 1][0] for j in xrange(number)] 197 198 self.axes.legend(lbars, self.legend, loc='best') 199 200 if self.title is not None: 201 self.axes.set_title(self.title) 202 203 def savefig(self, name): 204 self.figure.savefig(name) 205 206if __name__ == '__main__': 207 import random, sys 208 209 dim = 3 210 number = 5 211 212 args = sys.argv[1:] 213 if len(args) > 3: 214 sys.exit("invalid number of arguments") 215 elif len(args) > 0: 216 myshape = [ int(x) for x in args ] 217 else: 218 myshape = [ 3, 4, 8 ] 219 220 # generate a data matrix of the given shape 221 size = reduce(lambda x,y: x*y, myshape) 222 #data = [ random.randrange(size - i) + 10 for i in xrange(size) ] 223 data = [ float(i)/100.0 for i in xrange(size) ] 224 data = reshape(data, myshape) 225 226 # setup some test bar charts 227 if True: 228 chart1 = BarChart() 229 chart1.data = data 230 231 chart1.xlabel = 'Benchmark' 232 chart1.ylabel = 'Bandwidth (GBps)' 233 chart1.legend = [ 'x%d' % x for x in xrange(myshape[-1]) ] 234 chart1.xticks = [ 'xtick%d' % x for x in xrange(myshape[0]) ] 235 chart1.title = 'this is the title' 236 chart1.graph() 237 #chart1.savefig('/tmp/test1.png') 238 239 if False: 240 chart2 = BarChart() 241 chart2.data = data 242 chart2.colormap = 'gray' 243 chart2.graph() 244 #chart2.savefig('/tmp/test2.png') 245 246 pylab.show() 247