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