1
2import numpy as np
3import matplotlib.pyplot as plt
4
5import statsmodels.api as sm
6
7
8# Necessary to make horizontal axis labels fit
9plt.rcParams['figure.subplot.bottom'] = 0.23
10
11data = sm.datasets.anes96.load_pandas()
12party_ID = np.arange(7)
13labels = ["Strong Democrat", "Weak Democrat", "Independent-Democrat",
14          "Independent-Independent", "Independent-Republican",
15          "Weak Republican", "Strong Republican"]
16
17# Group age by party ID.
18age = [data.exog['age'][data.endog == id] for id in party_ID]
19
20
21# Create a violin plot.
22fig = plt.figure()
23ax = fig.add_subplot(111)
24
25sm.graphics.violinplot(age, ax=ax, labels=labels,
26                       plot_opts={'cutoff_val':5, 'cutoff_type':'abs',
27                                  'label_fontsize':'small',
28                                  'label_rotation':30})
29
30ax.set_xlabel("Party identification of respondent.")
31ax.set_ylabel("Age")
32ax.set_title("US national election '96 - Age & Party Identification")
33
34
35# Create a bean plot.
36fig2 = plt.figure()
37ax = fig2.add_subplot(111)
38
39sm.graphics.beanplot(age, ax=ax, labels=labels,
40                    plot_opts={'cutoff_val':5, 'cutoff_type':'abs',
41                               'label_fontsize':'small',
42                               'label_rotation':30})
43
44ax.set_xlabel("Party identification of respondent.")
45ax.set_ylabel("Age")
46ax.set_title("US national election '96 - Age & Party Identification")
47
48
49# Create a jitter plot.
50fig3 = plt.figure()
51ax = fig3.add_subplot(111)
52
53plot_opts={'cutoff_val':5, 'cutoff_type':'abs', 'label_fontsize':'small',
54           'label_rotation':30, 'violin_fc':(0.8, 0.8, 0.8),
55           'jitter_marker':'.', 'jitter_marker_size':3, 'bean_color':'#FF6F00',
56           'bean_mean_color':'#009D91'}
57sm.graphics.beanplot(age, ax=ax, labels=labels, jitter=True,
58                    plot_opts=plot_opts)
59
60ax.set_xlabel("Party identification of respondent.")
61ax.set_ylabel("Age")
62ax.set_title("US national election '96 - Age & Party Identification")
63
64
65# Create an asymmetrical jitter plot.
66ix = data.exog['income'] < 16  # incomes < $30k
67age = data.exog['age'][ix]
68endog = data.endog[ix]
69age_lower_income = [age[endog == id] for id in party_ID]
70
71ix = data.exog['income'] >= 20  # incomes > $50k
72age = data.exog['age'][ix]
73endog = data.endog[ix]
74age_higher_income = [age[endog == id] for id in party_ID]
75
76fig = plt.figure()
77ax = fig.add_subplot(111)
78
79plot_opts['violin_fc'] = (0.5, 0.5, 0.5)
80plot_opts['bean_show_mean'] = False
81plot_opts['bean_show_median'] = False
82plot_opts['bean_legend_text'] = r'Income < \$30k'
83plot_opts['cutoff_val'] = 10
84sm.graphics.beanplot(age_lower_income, ax=ax, labels=labels, side='left',
85                     jitter=True, plot_opts=plot_opts)
86plot_opts['violin_fc'] = (0.7, 0.7, 0.7)
87plot_opts['bean_color'] = '#009D91'
88plot_opts['bean_legend_text'] = r'Income > \$50k'
89sm.graphics.beanplot(age_higher_income, ax=ax, labels=labels, side='right',
90                     jitter=True, plot_opts=plot_opts)
91
92ax.set_xlabel("Party identification of respondent.")
93ax.set_ylabel("Age")
94ax.set_title("US national election '96 - Age & Party Identification")
95
96
97# Show all plots.
98plt.show()
99