barchart.py revision 1881:fc205a7edd58 
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