barchart.py revision 2118:1fe7d0ddf765 
1# Copyright (c) 2005-2006 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.font_manager import FontProperties
32from matplotlib.numerix import array, arange, reshape, shape, transpose, zeros
33from matplotlib.numerix import Float
34
35matplotlib.interactive(False)
36
37from chart import ChartOptions
38
39class BarChart(ChartOptions):
40    def __init__(self, default=None, **kwargs):
41        super(BarChart, self).__init__(default, **kwargs)
42        self.inputdata = None
43        self.chartdata = None
44
45    def gen_colors(self, count):
46        cmap = matplotlib.cm.get_cmap(self.colormap)
47        if count == 1:
48            return cmap([ 0.5 ])
49        else:
50            return cmap(arange(count) / float(count - 1))
51
52    # The input data format does not match the data format that the
53    # graph function takes because it is intuitive.  The conversion
54    # from input data format to chart data format depends on the
55    # dimensionality of the input data.  Check here for the
56    # dimensionality and correctness of the input data
57    def set_data(self, data):
58        if data is None:
59            self.inputdata = None
60            self.chartdata = None
61            return
62
63        data = array(data)
64        dim = len(shape(data))
65        if dim not in (1, 2, 3):
66            raise AttributeError, "Input data must be a 1, 2, or 3d matrix"
67        self.inputdata = data
68
69        # If the input data is a 1d matrix, then it describes a
70        # standard bar chart.
71        if dim == 1:
72            self.chartdata = array([[data]])
73
74        # If the input data is a 2d matrix, then it describes a bar
75        # chart with groups. The matrix being an array of groups of
76        # bars.
77        if dim == 2:
78            self.chartdata = transpose([data], axes=(2,0,1))
79
80        # If the input data is a 3d matrix, then it describes an array
81        # of groups of bars with each bar being an array of stacked
82        # values.
83        if dim == 3:
84            self.chartdata = transpose(data, axes=(1,2,0))
85
86    def get_data(self):
87        return self.inputdata
88
89    data = property(get_data, set_data)
90
91    # Graph the chart data.
92    # Input is a 3d matrix that describes a plot that has multiple
93    # groups, multiple bars in each group, and multiple values stacked
94    # in each bar.  The underlying bar() function expects a sequence of
95    # bars in the same stack location and same group location, so the
96    # organization of the matrix is that the inner most sequence
97    # represents one of these bar groups, then those are grouped
98    # together to make one full stack of bars in each group, and then
99    # the outer most layer describes the groups.  Here is an example
100    # data set and how it gets plotted as a result.
101    #
102    # e.g. data = [[[10,11,12], [13,14,15],  [16,17,18], [19,20,21]],
103    #              [[22,23,24], [25,26,27],  [28,29,30], [31,32,33]]]
104    #
105    # will plot like this:
106    #
107    #    19 31    20 32    21 33
108    #    16 28    17 29    18 30
109    #    13 25    14 26    15 27
110    #    10 22    11 23    12 24
111    #
112    # Because this arrangement is rather conterintuitive, the rearrange
113    # function takes various matricies and arranges them to fit this
114    # profile.
115    #
116    # This code deals with one of the dimensions in the matrix being
117    # one wide.
118    #
119    def graph(self):
120        if self.chartdata is None:
121            raise AttributeError, "Data not set for bar chart!"
122
123        need_subticks = True
124
125        dim = len(shape(self.inputdata))
126        cshape = shape(self.chartdata)
127        print cshape
128        if dim == 1:
129            colors = self.gen_colors(cshape[2])
130            colors = [ [ colors ] * cshape[1] ] * cshape[0]
131            need_subticks = False
132
133        if dim == 2:
134            colors = self.gen_colors(cshape[0])
135            colors = [ [ [ c ] * cshape[2] ] * cshape[1] for c in colors ]
136
137        if dim == 3:
138            colors = self.gen_colors(cshape[1])
139            colors = [ [ [ c ] * cshape[2] for c in colors ] ] * cshape[0]
140
141        colors = array(colors)
142
143        self.figure = pylab.figure(figsize=self.chart_size)
144
145        outer_axes = None
146        inner_axes = None
147        if need_subticks:
148            self.metaaxes = self.figure.add_axes(self.figure_size)
149            self.metaaxes.set_yticklabels([])
150            self.metaaxes.set_yticks([])
151            size = [0] * 4
152            size[0] = self.figure_size[0]
153            size[1] = self.figure_size[1] + .075
154            size[2] = self.figure_size[2]
155            size[3] = self.figure_size[3] - .075
156            self.axes = self.figure.add_axes(size)
157            outer_axes = self.metaaxes
158            inner_axes = self.axes
159        else:
160            self.axes = self.figure.add_axes(self.figure_size)
161            outer_axes = self.axes
162            inner_axes = self.axes
163
164        bars_in_group = len(self.chartdata)
165        if bars_in_group < 5:
166            width = 1.0 / ( bars_in_group + 1)
167            center = width / 2
168        else:
169            width = .8 / bars_in_group
170            center = .1
171
172        bars = []
173        for i,stackdata in enumerate(self.chartdata):
174            bottom = array([0.0] * len(stackdata[0]), Float)
175            stack = []
176            for j,bardata in enumerate(stackdata):
177                bardata = array(bardata)
178                ind = arange(len(bardata)) + i * width + center
179                bar = self.axes.bar(ind, bardata, width, bottom=bottom,
180                                    color=colors[i][j])
181                if dim != 1:
182                    self.metaaxes.bar(ind, [0] * len(bardata), width)
183                stack.append(bar)
184                bottom += bardata
185            bars.append(stack)
186
187        if self.xlabel is not None:
188            outer_axes.set_xlabel(self.xlabel)
189
190        if self.ylabel is not None:
191            inner_axes.set_ylabel(self.ylabel)
192
193        if self.yticks is not None:
194            ymin, ymax = self.axes.get_ylim()
195            nticks = float(len(self.yticks))
196            ticks = arange(nticks) / (nticks - 1) * (ymax - ymin)  + ymin
197            inner_axes.set_yticks(ticks)
198            inner_axes.set_yticklabels(self.yticks)
199        elif self.ylim is not None:
200            self.inner_axes.set_ylim(self.ylim)
201
202        if self.xticks is not None:
203            outer_axes.set_xticks(arange(cshape[2]) + .5)
204            outer_axes.set_xticklabels(self.xticks)
205        if self.xsubticks is not None:
206            inner_axes.set_xticks(arange((cshape[0] + 1)*cshape[2])*width + 2*center)
207            self.xsubticks.append('')
208            inner_axes.set_xticklabels(self.xsubticks * cshape[2], fontsize=8)
209        if self.legend is not None:
210            if dim == 1:
211                lbars = bars[0][0]
212            if dim == 2:
213                lbars = [ bars[i][0][0] for i in xrange(len(bars))]
214            if dim == 3:
215                number = len(bars[0])
216                lbars = [ bars[0][number - j - 1][0] for j in xrange(number)]
217
218            self.figure.legend(lbars, self.legend, self.legend_loc,
219                               prop=FontProperties(size=self.legend_size))
220
221        if self.title is not None:
222            self.axes.set_title(self.title)
223
224    def savefig(self, name):
225        self.figure.savefig(name)
226
227    def savecsv(self, name):
228        f = file(name, 'w')
229        data = array(self.inputdata)
230        dim = len(data.shape)
231
232        if dim == 1:
233            #if self.xlabel:
234            #    f.write(', '.join(list(self.xlabel)) + '\n')
235            f.write(', '.join([ '%f' % val for val in data]) + '\n')
236        if dim == 2:
237            #if self.xlabel:
238            #    f.write(', '.join([''] + list(self.xlabel)) + '\n')
239            for i,row in enumerate(data):
240                ylabel = []
241                #if self.ylabel:
242                #    ylabel = [ self.ylabel[i] ]
243                f.write(', '.join(ylabel + [ '%f' % val for val in row]) + '\n')
244        if dim == 3:
245            f.write("don't do 3D csv files\n")
246            pass
247
248        f.close()
249
250if __name__ == '__main__':
251    from random import randrange
252    import random, sys
253
254    dim = 3
255    number = 5
256
257    args = sys.argv[1:]
258    if len(args) > 3:
259        sys.exit("invalid number of arguments")
260    elif len(args) > 0:
261        myshape = [ int(x) for x in args ]
262    else:
263        myshape = [ 3, 4, 8 ]
264
265    # generate a data matrix of the given shape
266    size = reduce(lambda x,y: x*y, myshape)
267    #data = [ random.randrange(size - i) + 10 for i in xrange(size) ]
268    data = [ float(i)/100.0 for i in xrange(size) ]
269    data = reshape(data, myshape)
270
271    # setup some test bar charts
272    if True:
273        chart1 = BarChart()
274        chart1.data = data
275
276        chart1.xlabel = 'Benchmark'
277        chart1.ylabel = 'Bandwidth (GBps)'
278        chart1.legend = [ 'x%d' % x for x in xrange(myshape[-1]) ]
279        chart1.xticks = [ 'xtick%d' % x for x in xrange(myshape[0]) ]
280        chart1.title = 'this is the title'
281        chart1.figure_size = [0.1, 0.2, 0.7, 0.85 ]
282        if len(myshape) > 1:
283            chart1.xsubticks = [ '%d' % x for x in xrange(myshape[1]) ]
284        chart1.graph()
285        chart1.savefig('/tmp/test1.png')
286        chart1.savefig('/tmp/test1.ps')
287        chart1.savefig('/tmp/test1.eps')
288        chart1.savecsv('/tmp/test1.csv')
289
290    if False:
291        chart2 = BarChart()
292        chart2.data = data
293        chart2.colormap = 'gray'
294        chart2.graph()
295        chart2.savefig('/tmp/test2.png')
296        chart2.savefig('/tmp/test2.ps')
297
298    pylab.myshow()
299