forked from statsmodels/statsmodels
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
149 lines (120 loc) · 4.04 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
"""Helper functions for graphics with Matplotlib."""
from statsmodels.compat.python import lrange, range
__all__ = ['create_mpl_ax', 'create_mpl_fig']
def _import_mpl():
"""This function is not needed outside this utils module."""
try:
import matplotlib.pyplot as plt
except:
raise ImportError("Matplotlib is not found.")
return plt
def create_mpl_ax(ax=None):
"""Helper function for when a single plot axis is needed.
Parameters
----------
ax : Matplotlib AxesSubplot instance, optional
If given, this subplot is used to plot in instead of a new figure being
created.
Returns
-------
fig : Matplotlib figure instance
If `ax` is None, the created figure. Otherwise the figure to which
`ax` is connected.
ax : Matplotlib AxesSubplot instance
The created axis if `ax` is None, otherwise the axis that was passed
in.
Notes
-----
This function imports `matplotlib.pyplot`, which should only be done to
create (a) figure(s) with ``plt.figure``. All other functionality exposed
by the pyplot module can and should be imported directly from its
Matplotlib module.
See Also
--------
create_mpl_fig
Examples
--------
A plotting function has a keyword ``ax=None``. Then calls:
>>> from statsmodels.graphics import utils
>>> fig, ax = utils.create_mpl_ax(ax)
"""
if ax is None:
plt = _import_mpl()
fig = plt.figure()
ax = fig.add_subplot(111)
else:
fig = ax.figure
return fig, ax
def create_mpl_fig(fig=None, figsize=None):
"""Helper function for when multiple plot axes are needed.
Those axes should be created in the functions they are used in, with
``fig.add_subplot()``.
Parameters
----------
fig : Matplotlib figure instance, optional
If given, this figure is simply returned. Otherwise a new figure is
created.
Returns
-------
fig : Matplotlib figure instance
If `fig` is None, the created figure. Otherwise the input `fig` is
returned.
See Also
--------
create_mpl_ax
"""
if fig is None:
plt = _import_mpl()
fig = plt.figure(figsize=figsize)
return fig
def maybe_name_or_idx(idx, model):
"""
Give a name or an integer and return the name and integer location of the
column in a design matrix.
"""
if idx is None:
idx = lrange(model.exog.shape[1])
if isinstance(idx, int):
exog_name = model.exog_names[idx]
exog_idx = idx
# anticipate index as list and recurse
elif isinstance(idx, (tuple, list)):
exog_name = []
exog_idx = []
for item in idx:
exog_name_item, exog_idx_item = maybe_name_or_idx(item, model)
exog_name.append(exog_name_item)
exog_idx.append(exog_idx_item)
else: # assume we've got a string variable
exog_name = idx
exog_idx = model.exog_names.index(idx)
return exog_name, exog_idx
def get_data_names(series_or_dataframe):
"""
Input can be an array or pandas-like. Will handle 1d array-like but not
2d. Returns a str for 1d data or a list of strings for 2d data.
"""
names = getattr(series_or_dataframe, 'name', None)
if not names:
names = getattr(series_or_dataframe, 'columns', None)
if not names:
shape = getattr(series_or_dataframe, 'shape', [1])
nvars = 1 if len(shape) == 1 else series.shape[1]
names = ["X%d" for names in range(nvars)]
if nvars == 1:
names = names[0]
else:
names = names.tolist()
return names
def annotate_axes(index, labels, points, offset_points, size, ax, **kwargs):
"""
Annotate Axes with labels, points, offset_points according to the
given index.
"""
for i in index:
label = labels[i]
point = points[i]
offset = offset_points[i]
ax.annotate(label, point, xytext=offset, textcoords="offset points",
size=size, **kwargs)
return ax