1import numpy as np
2import pytest
3
4from statsmodels.datasets import anes96
5from statsmodels.graphics.boxplots import beanplot, violinplot
6
7try:
8    import matplotlib.pyplot as plt
9except ImportError:
10    pass
11
12
13@pytest.fixture(scope="module")
14def age_and_labels():
15    # Test violinplot and beanplot with the same dataset.
16    data = anes96.load_pandas()
17    party_ID = np.arange(7)
18    labels = ["Strong Democrat", "Weak Democrat", "Independent-Democrat",
19              "Independent-Independent", "Independent-Republican",
20              "Weak Republican", "Strong Republican"]
21    age = [data.exog['age'][data.endog == id] for id in party_ID]
22    age = np.array(age, dtype="object")
23    return age, labels
24
25
26# TODO: Remove once SciPy 1.5.0 is the minimum
27IGNORE_VISIBLE_DEPR = """\
28ignore:Creating an ndarray from ragged nested sequences:\
29numpy.VisibleDeprecationWarning:"""
30
31
32@pytest.mark.filterwarnings(IGNORE_VISIBLE_DEPR)
33@pytest.mark.matplotlib
34def test_violinplot(age_and_labels, close_figures):
35    age, labels = age_and_labels
36
37    fig, ax = plt.subplots(1, 1)
38    violinplot(age, ax=ax, labels=labels,
39               plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
40                          'label_fontsize': 'small',
41                          'label_rotation': 30})
42
43
44@pytest.mark.filterwarnings(IGNORE_VISIBLE_DEPR)
45@pytest.mark.matplotlib
46def test_violinplot_bw_factor(age_and_labels, close_figures):
47    age, labels = age_and_labels
48
49    fig, ax = plt.subplots(1, 1)
50    violinplot(age, ax=ax, labels=labels,
51               plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
52                          'label_fontsize': 'small',
53                          'label_rotation': 30,
54                          'bw_factor': .2})
55
56
57@pytest.mark.matplotlib
58def test_beanplot(age_and_labels, close_figures):
59    age, labels = age_and_labels
60
61    fig, ax = plt.subplots(1, 1)
62    beanplot(age, ax=ax, labels=labels,
63             plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
64                        'label_fontsize': 'small',
65                        'label_rotation': 30})
66
67
68@pytest.mark.matplotlib
69def test_beanplot_jitter(age_and_labels, close_figures):
70    age, labels = age_and_labels
71
72    fig, ax = plt.subplots(1, 1)
73    beanplot(age, ax=ax, labels=labels, jitter=True,
74             plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
75                        'label_fontsize': 'small',
76                        'label_rotation': 30})
77
78
79@pytest.mark.matplotlib
80def test_beanplot_side_right(age_and_labels, close_figures):
81    age, labels = age_and_labels
82
83    fig, ax = plt.subplots(1, 1)
84    beanplot(age, ax=ax, labels=labels, jitter=True, side='right',
85             plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
86                        'label_fontsize': 'small',
87                        'label_rotation': 30})
88
89
90@pytest.mark.matplotlib
91def test_beanplot_side_left(age_and_labels, close_figures):
92    age, labels = age_and_labels
93
94    fig, ax = plt.subplots(1, 1)
95    beanplot(age, ax=ax, labels=labels, jitter=True, side='left',
96             plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
97                        'label_fontsize': 'small',
98                        'label_rotation': 30})
99
100
101@pytest.mark.matplotlib
102def test_beanplot_legend_text(age_and_labels, close_figures):
103    age, labels = age_and_labels
104
105    fig, ax = plt.subplots(1, 1)
106    beanplot(age, ax=ax, labels=labels,
107             plot_opts={'bean_legend_text': 'text'})
108