1# -*- coding: utf-8 -*-
2"""
3Created on Fri May 30 22:56:57 2014
4
5Author: Josef Perktold
6License: BSD-3
7
8"""
9
10import numpy as np
11from numpy.testing import assert_allclose, assert_raises
12
13from statsmodels.base._constraints import (
14    TransformRestriction,
15    fit_constrained,
16    transform_params_constraint,
17)
18
19if __name__ == '__main__':
20
21
22
23    R = np.array([[1, 1, 0, 0, 0], [0, 0, 1, -1, 0]])
24
25
26
27    k_constr, k_vars = R.shape
28
29    m = np.eye(k_vars) - R.T.dot(np.linalg.pinv(R).T)
30    evals, evecs = np.linalg.eigh(m)
31
32    L = evecs[:, :k_constr]
33    T = evecs[:, k_constr:]
34
35    print(T.T.dot(np.eye(k_vars)))
36
37    tr = np.column_stack((T, R.T))
38
39    q = [2, 0]
40    tr0 = TransformRestriction(R, q)
41
42    p_reduced = [1,1,1]
43    #round trip test
44    assert_allclose(tr0.reduce(tr0.expand(p_reduced)), p_reduced, rtol=1e-14)
45
46
47    p = tr0.expand(p_reduced)
48    assert_allclose(R.dot(p), q, rtol=1e-14)
49
50    # inconsistent restrictions
51    #Ri = np.array([[1, 1, 0, 0, 0], [0, 0, 1, -1, 0], [0, 0, 1, -2, 0]])
52    R = np.array([[1, 1, 0, 0, 0], [0, 0, 1, -1, 0], [0, 0, 1, 0, -1]])
53    q = np.zeros(R.shape[0])
54    tr1 = TransformRestriction(R, q)  # bug raises error with q
55    p = tr1.expand([1,1])
56
57    # inconsistent restrictions that has a solution with p3=0
58    Ri = np.array([[1, 1, 0, 0, 0], [0, 0, 1, -1, 0], [0, 0, 1, -2, 0]])
59    tri = TransformRestriction(Ri, [0, 1, 1])
60    # Note: the only way this can hold is if variable 3 is zero
61    p = tri.expand([1,1])
62    print(p[[2,3]])
63    #array([  1.00000000e+00,  -1.34692639e-17])
64
65    # inconsistent without any possible solution
66    Ri2 = np.array([[0, 0, 0, 1, 0], [0, 0, 1, -1, 0], [0, 0, 1, -2, 0]])
67    q = [1, 1]
68    #tri2 = TransformRestriction(Ri2, q)
69    #p = tri.expand([1,1])
70    assert_raises(ValueError, TransformRestriction, Ri2, q)
71    # L does not have full row rank, calculating constant fails with Singular Matrix
72
73    # transform data xr = T x
74    np.random.seed(1)
75    x = np.random.randn(10, 5)
76    xr = tr1.reduce(x)
77    # roundtrip
78    x2 = tr1.expand(xr)
79    # this does not hold ? do not use constant? do not need it anyway ?
80    #assert_allclose(x2, x, rtol=1e-14)
81
82
83    from patsy import DesignInfo
84
85    names = 'a b c d'.split()
86    LC = DesignInfo(names).linear_constraint('a + b = 0')
87    LC = DesignInfo(names).linear_constraint(['a + b = 0', 'a + 2*c = 1', 'b-a', 'c-a', 'd-a'])
88    #LC = DesignInfo(self.model.exog_names).linear_constraint(r_matrix)
89    r_matrix, q_matrix = LC.coefs, LC.constants
90
91    np.random.seed(123)
92    nobs = 20
93    x = 1 + np.random.randn(nobs, 4)
94    exog = np.column_stack((np.ones(nobs), x))
95    endog = exog.sum(1) + np.random.randn(nobs)
96
97    from statsmodels.regression.linear_model import OLS
98    res2 = OLS(endog, exog).fit()
99    #transf = TransformRestriction(np.eye(exog.shape[1])[:2], res2.params[:2] / 2)
100    transf = TransformRestriction([[0, 0, 0,1,1]], res2.params[-2:].sum())
101    exog_st = transf.reduce(exog)
102    res1 = OLS(endog, exog_st).fit()
103    # need to correct for constant/offset in the optimization
104    res1 = OLS(endog - exog.dot(transf.constant.squeeze()), exog_st).fit()
105    params = transf.expand(res1.params).squeeze()
106    assert_allclose(params, res2.params, rtol=1e-13)
107    print(res2.params)
108    print(params)
109    print(res1.params)
110
111    res3_ols = OLS(endog - exog[:, -1], exog[:, :-2]).fit()
112    #transf = TransformRestriction(np.eye(exog.shape[1])[:2], res2.params[:2] / 2)
113    transf3 = TransformRestriction([[0, 0, 0, 1, 0],[0, 0, 0, 0, 1]], [0, 1])
114    exog3_st = transf3.reduce(exog)
115    res3 = OLS(endog, exog3_st).fit()
116    # need to correct for constant/offset in the optimization
117    res3 = OLS(endog - exog.dot(transf3.constant.squeeze()), exog3_st).fit()
118    params = transf3.expand(res3.params).squeeze()
119    assert_allclose(params[:-2], res3_ols.params, rtol=1e-13)
120    print(res3.params)
121    print(params)
122    print(res3_ols.params)
123    print(res3_ols.bse)
124    # the following raises `ValueError: cannot test a constant constraint`
125    #tt = res3.t_test(transf3.transf_mat, transf3.constant.squeeze())
126    #print tt.sd
127    cov_params3 = transf3.transf_mat.dot(res3.cov_params()).dot(transf3.transf_mat.T)
128    bse3 = np.sqrt(np.diag(cov_params3))
129    print(bse3)
130
131    tp = transform_params_constraint(res2.params, res2.normalized_cov_params,
132                                     transf3.R, transf3.q)
133    tp = transform_params_constraint(res2.params, res2.cov_params(), transf3.R, transf3.q)
134
135    import statsmodels.api as sm
136    rand_data = sm.datasets.randhie.load()
137    rand_exog = rand_data.exog.view(float).reshape(len(rand_data.exog), -1)
138    rand_exog = sm.add_constant(rand_exog, prepend=False)
139
140
141    # Fit Poisson model:
142    poisson_mod0 = sm.Poisson(rand_data.endog, rand_exog)
143    poisson_res0 = poisson_mod0.fit(method="newton")
144
145    R = np.zeros((2, 10))
146    R[0, -2] = 1
147    R[1, -1] = 1
148    transfp = TransformRestriction(R, [0, 1])
149    poisson_mod = sm.Poisson(rand_data.endog, rand_exog[:, :-2])
150    # note wrong offset, why did I put offset in fit ? it's ignored
151    poisson_res = poisson_mod.fit(method="newton", offset=rand_exog.dot(transfp.constant.squeeze()))
152
153    exogp_st = transfp.reduce(rand_exog)
154    poisson_modr = sm.Poisson(rand_data.endog, exogp_st)
155    poisson_resr = poisson_modr.fit(method="newton")
156    paramsp = transfp.expand(poisson_resr.params).squeeze()
157    print('\nPoisson')
158    print(paramsp)
159    print(poisson_res.params)
160    # error because I do not use the unconstrained basic model
161#    tp = transform_params_constraint(poisson_res.params, poisson_res.cov_params(), transfp.R, transfp.q)
162#    cov_params3 = transf3.transf_mat.dot(res3.cov_params()).dot(transf3.transf_mat.T)
163#    bse3 = np.sqrt(np.diag(cov_params3))
164
165
166    poisson_mod0 = sm.Poisson(rand_data.endog, rand_exog)
167    poisson_res0 = poisson_mod0.fit(method="newton")
168    tp = transform_params_constraint(poisson_res0.params, poisson_res0.cov_params(), transfp.R, transfp.q)
169    cov_params3 = transf3.transf_mat.dot(res3.cov_params()).dot(transf3.transf_mat.T)
170    bse3 = np.sqrt(np.diag(cov_params3))
171
172    # try again same example as it was intended
173
174    poisson_mod = sm.Poisson(rand_data.endog, rand_exog[:, :-2], offset=rand_exog[:, -1])
175    poisson_res = poisson_mod.fit(method="newton")
176
177    exogp_st = transfp.reduce(rand_exog)
178    poisson_modr = sm.Poisson(rand_data.endog, exogp_st, offset=rand_exog.dot(transfp.constant.squeeze()))
179    poisson_resr = poisson_modr.fit(method="newton")
180    paramsp = transfp.expand(poisson_resr.params).squeeze()
181    print('\nPoisson')
182    print(paramsp)
183    print(poisson_resr.params)
184    tp = transform_params_constraint(poisson_res0.params, poisson_res0.cov_params(), transfp.R, transfp.q)
185    cov_paramsp = transfp.transf_mat.dot(poisson_resr.cov_params()).dot(transfp.transf_mat.T)
186    bsep = np.sqrt(np.diag(cov_paramsp))
187    print(bsep)
188    p, cov, res_r = fit_constrained(poisson_mod0, transfp.R, transfp.q)
189    se = np.sqrt(np.diag(cov))
190    print(p)
191    print(se)
192