100 lines
3.1 KiB
Python
100 lines
3.1 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from statsmodels.datasets import anes96
|
|
from statsmodels.graphics.boxplots import beanplot, violinplot
|
|
|
|
try:
|
|
import matplotlib.pyplot as plt
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def age_and_labels():
|
|
# Test violinplot and beanplot with the same dataset.
|
|
data = anes96.load_pandas()
|
|
party_ID = np.arange(7)
|
|
labels = ["Strong Democrat", "Weak Democrat", "Independent-Democrat",
|
|
"Independent-Independent", "Independent-Republican",
|
|
"Weak Republican", "Strong Republican"]
|
|
age = [data.exog['age'][data.endog == id] for id in party_ID]
|
|
age = np.array(age, dtype="object")
|
|
return age, labels
|
|
|
|
|
|
@pytest.mark.matplotlib
|
|
def test_violinplot(age_and_labels, close_figures):
|
|
age, labels = age_and_labels
|
|
|
|
fig, ax = plt.subplots(1, 1)
|
|
violinplot(age, ax=ax, labels=labels,
|
|
plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
|
|
'label_fontsize': 'small',
|
|
'label_rotation': 30})
|
|
|
|
|
|
@pytest.mark.matplotlib
|
|
def test_violinplot_bw_factor(age_and_labels, close_figures):
|
|
age, labels = age_and_labels
|
|
|
|
fig, ax = plt.subplots(1, 1)
|
|
violinplot(age, ax=ax, labels=labels,
|
|
plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
|
|
'label_fontsize': 'small',
|
|
'label_rotation': 30,
|
|
'bw_factor': .2})
|
|
|
|
|
|
@pytest.mark.matplotlib
|
|
def test_beanplot(age_and_labels, close_figures):
|
|
age, labels = age_and_labels
|
|
|
|
fig, ax = plt.subplots(1, 1)
|
|
beanplot(age, ax=ax, labels=labels,
|
|
plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
|
|
'label_fontsize': 'small',
|
|
'label_rotation': 30})
|
|
|
|
|
|
@pytest.mark.matplotlib
|
|
def test_beanplot_jitter(age_and_labels, close_figures):
|
|
age, labels = age_and_labels
|
|
|
|
fig, ax = plt.subplots(1, 1)
|
|
beanplot(age, ax=ax, labels=labels, jitter=True,
|
|
plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
|
|
'label_fontsize': 'small',
|
|
'label_rotation': 30})
|
|
|
|
|
|
@pytest.mark.matplotlib
|
|
def test_beanplot_side_right(age_and_labels, close_figures):
|
|
age, labels = age_and_labels
|
|
|
|
fig, ax = plt.subplots(1, 1)
|
|
beanplot(age, ax=ax, labels=labels, jitter=True, side='right',
|
|
plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
|
|
'label_fontsize': 'small',
|
|
'label_rotation': 30})
|
|
|
|
|
|
@pytest.mark.matplotlib
|
|
def test_beanplot_side_left(age_and_labels, close_figures):
|
|
age, labels = age_and_labels
|
|
|
|
fig, ax = plt.subplots(1, 1)
|
|
beanplot(age, ax=ax, labels=labels, jitter=True, side='left',
|
|
plot_opts={'cutoff_val': 5, 'cutoff_type': 'abs',
|
|
'label_fontsize': 'small',
|
|
'label_rotation': 30})
|
|
|
|
|
|
@pytest.mark.matplotlib
|
|
def test_beanplot_legend_text(age_and_labels, close_figures):
|
|
age, labels = age_and_labels
|
|
|
|
fig, ax = plt.subplots(1, 1)
|
|
beanplot(age, ax=ax, labels=labels,
|
|
plot_opts={'bean_legend_text': 'text'})
|