forked from statsmodels/statsmodels
-
Notifications
You must be signed in to change notification settings - Fork 0
/
wrapper.py
145 lines (114 loc) · 4.04 KB
/
wrapper.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
import inspect
import functools
import types
import numpy as np
class ResultsWrapper(object):
"""
Class which wraps a statsmodels estimation Results class and steps in to
reattach metadata to results (if available)
"""
_wrap_attrs = {}
_wrap_methods = {}
def __init__(self, results):
self._results = results
self.__doc__ = results.__doc__
def __dir__(self):
return [x for x in dir(self._results)]
def __getattribute__(self, attr):
get = lambda name: object.__getattribute__(self, name)
try:
results = get('_results')
except AttributeError:
pass
try:
return get(attr)
except AttributeError:
pass
obj = getattr(results, attr)
data = results.model._data
how = self._wrap_attrs.get(attr)
if how:
obj = data.wrap_output(obj, how=how)
return obj
def __getstate__(self):
#print 'pickling wrapper', self.__dict__
return self.__dict__
def __setstate__(self, dict_):
#print 'unpickling wrapper', dict_
self.__dict__.update(dict_)
def save(self, fname, remove_data=False):
'''save a pickle of this instance
Parameters
----------
fname : string or filehandle
fname can be a string to a file path or filename, or a filehandle.
remove_data : bool
If False (default), then the instance is pickled without changes.
If True, then all arrays with length nobs are set to None before
pickling. See the remove_data method.
In some cases not all arrays will be set to None.
'''
from statsmodels.iolib.smpickle import save_pickle
if remove_data:
self.remove_data()
save_pickle(self, fname)
@classmethod
def load(cls, fname):
from statsmodels.iolib.smpickle import load_pickle
return load_pickle(fname)
def union_dicts(*dicts):
result = {}
for d in dicts:
result.update(d)
return result
def make_wrapper(func, how):
@functools.wraps(func)
def wrapper(self, *args, **kwargs):
results = object.__getattribute__(self, '_results')
data = results.model._data
return data.wrap_output(func(results, *args, **kwargs), how)
argspec = inspect.getargspec(func)
formatted = inspect.formatargspec(argspec[0], varargs=argspec[1],
defaults=argspec[3])
try:
func_name = func.im_func.func_name
except AttributeError:
#Python 3
func_name = func.__name__
wrapper.__doc__ = "%s%s\n%s" % (func_name, formatted, wrapper.__doc__)
return wrapper
def populate_wrapper(klass, wrapping):
for meth, how in klass._wrap_methods.iteritems():
if not hasattr(wrapping, meth):
continue
func = getattr(wrapping, meth)
wrapper = make_wrapper(func, how)
setattr(klass, meth, wrapper)
if __name__ == '__main__':
import statsmodels.api as sm
from pandas import DataFrame
data = sm.datasets.longley.load()
df = DataFrame(data.exog, columns=data.exog_name)
y = data.endog
# data.exog = sm.add_constant(data.exog)
df['intercept'] = 1.
olsresult = sm.OLS(y, df).fit()
rlmresult = sm.RLM(y, df).fit()
# olswrap = RegressionResultsWrapper(olsresult)
# rlmwrap = RLMResultsWrapper(rlmresult)
data = sm.datasets.wfs.load()
# get offset
offset = np.log(data.exog[:,-1])
exog = data.exog[:,:-1]
# convert dur to dummy
exog = sm.tools.categorical(exog, col=0, drop=True)
# drop reference category
# convert res to dummy
exog = sm.tools.categorical(exog, col=0, drop=True)
# convert edu to dummy
exog = sm.tools.categorical(exog, col=0, drop=True)
# drop reference categories and add intercept
exog = sm.add_constant(exog[:,[1,2,3,4,5,7,8,10,11,12]])
endog = np.round(data.endog)
mod = sm.GLM(endog, exog, family=sm.families.Poisson()).fit()
# glmwrap = GLMResultsWrapper(mod)