|
| 1 | +""" |
| 2 | +============================================= |
| 3 | +Discrete distribution as horizontal bar chart |
| 4 | +============================================= |
| 5 | +
|
| 6 | +Stacked bar charts can be used to visualize discrete distributions. |
| 7 | +
|
| 8 | +This example visualizes the result of a survey in which people could rate |
| 9 | +their agreement to questions on a five-element scale. |
| 10 | +
|
| 11 | +The horizontal stacking is achieved by calling `~.Axes.barh()` for each |
| 12 | +category and passing the starting point as the cumulative sum of the |
| 13 | +already drawn bars via the parameter ``left``. |
| 14 | +""" |
| 15 | + |
| 16 | +import numpy as np |
| 17 | +import matplotlib.pyplot as plt |
| 18 | + |
| 19 | + |
| 20 | +category_names = ['Strongly disagree', 'Disagree', |
| 21 | + 'Neither agree nor disagree', 'Agree', 'Strongly agree'] |
| 22 | +results = { |
| 23 | + 'Question 1': [10, 15, 17, 32, 26], |
| 24 | + 'Question 2': [26, 22, 29, 10, 13], |
| 25 | + 'Question 3': [35, 37, 7, 2, 19], |
| 26 | + 'Question 4': [32, 11, 9, 15, 33], |
| 27 | + 'Question 5': [21, 29, 5, 5, 40], |
| 28 | + 'Question 6': [8, 19, 5, 30, 38] |
| 29 | +} |
| 30 | + |
| 31 | + |
| 32 | +def survey(results, category_names): |
| 33 | + """ |
| 34 | + Parameters |
| 35 | + ---------- |
| 36 | + results : dict |
| 37 | + A mapping from question labels to a list of answers per category. |
| 38 | + It is assumed all lists contain the same number of entries and that |
| 39 | + it matches the length of *category_names*. |
| 40 | + category_names : list of str |
| 41 | + The category labels. |
| 42 | + """ |
| 43 | + labels = list(results.keys()) |
| 44 | + data = np.array(list(results.values())) |
| 45 | + data_cum = data.cumsum(axis=1) |
| 46 | + category_colors = plt.get_cmap('RdYlGn')( |
| 47 | + np.linspace(0.15, 0.85, data.shape[1])) |
| 48 | + |
| 49 | + fig, ax = plt.subplots(figsize=(9.2, 5)) |
| 50 | + ax.invert_yaxis() |
| 51 | + ax.xaxis.set_visible(False) |
| 52 | + ax.set_xlim(0, np.sum(data, axis=1).max()) |
| 53 | + |
| 54 | + for i, (colname, color) in enumerate(zip(category_names, category_colors)): |
| 55 | + widths = data[:, i] |
| 56 | + starts = data_cum[:, i] - widths |
| 57 | + ax.barh(labels, widths, left=starts, height=0.5, |
| 58 | + label=colname, color=color) |
| 59 | + xcenters = starts + widths / 2 |
| 60 | + |
| 61 | + r, g, b, _ = color |
| 62 | + text_color = 'white' if r * g * b < 0.5 else 'darkgrey' |
| 63 | + for y, (x, c) in enumerate(zip(xcenters, widths)): |
| 64 | + ax.text(x, y, str(int(c)), ha='center', va='center', |
| 65 | + color=text_color) |
| 66 | + ax.legend(ncol=len(category_names), bbox_to_anchor=(0, 1), |
| 67 | + loc='lower left', fontsize='small') |
| 68 | + |
| 69 | + return fig, ax |
| 70 | + |
| 71 | + |
| 72 | +survey(results, category_names) |
| 73 | + |
| 74 | +############################################################################# |
| 75 | +# |
| 76 | +# ------------ |
| 77 | +# |
| 78 | +# References |
| 79 | +# """""""""" |
| 80 | +# |
| 81 | +# The use of the following functions, methods, classes and modules is shown |
| 82 | +# in this example: |
| 83 | + |
| 84 | +import matplotlib |
| 85 | +matplotlib.axes.Axes.barh |
| 86 | +matplotlib.pyplot.barh |
| 87 | +matplotlib.axes.Axes.text |
| 88 | +matplotlib.pyplot.text |
| 89 | +matplotlib.axes.Axes.legend |
| 90 | +matplotlib.pyplot.legend |
0 commit comments