BUG/ENH: modify stationary cov_struct for GEE

statsmodels/genmod/cov_struct.py

@@ -131,7 +131,7 @@ def covariance_matrix_solve(self, expval, index, stdev, rhs):
 
         If the covariance matrix is singular or not SPD, it is
         projected to the nearest such matrix.  These projection events
-        are recorded in the fit_history member of the GEE model.
+        are recorded in the fit_history attribute of the GEE model.
 
         Systems of linear equations with the covariance matrix as the
         left hand side (LHS) are solved for different right hand sides
@@ -161,6 +161,9 @@ def covariance_matrix_solve(self, expval, index, stdev, rhs):
                                    threshold=threshold)
                 threshold *= 2
                 cov_adjust += 1
+                msg = "At least one covariance matrix was not PSD "
+                msg += "and required projection."
+                warnings.warn(msg)
 
         self.cov_adjust.append(cov_adjust)
 
@@ -524,7 +527,7 @@ def __init__(self, max_lag=1, grid=False):
         super(Stationary, self).__init__()
         self.max_lag = max_lag
         self.grid = grid
-        self.dep_params = np.zeros(max_lag)
+        self.dep_params = np.zeros(max_lag + 1)
 
     def initialize(self, model):
 
@@ -557,10 +560,11 @@ def update_grid(self, params):
 
             dep_params[0] += np.sum(resid * resid) / len(resid)
             for j in range(1, self.max_lag + 1):
-                dep_params[j] += np.sum(resid[0:-j] *
-                                        resid[j:]) / len(resid[j:])
+                v = resid[j:]
+                dep_params[j] += np.sum(resid[0:-j] * v) / len(v)
 
-        self.dep_params = dep_params[1:] / dep_params[0]
+        dep_params /= dep_params[0]
+        self.dep_params = dep_params
 
     def update_nogrid(self, params):
 
@@ -570,55 +574,63 @@ def update_nogrid(self, params):
 
         dep_params = np.zeros(self.max_lag + 1)
         dn = np.zeros(self.max_lag + 1)
+        resid_ssq = 0
+        resid_ssq_n = 0
         for i in range(self.model.num_group):
 
             expval, _ = cached_means[i]
             stdev = np.sqrt(varfunc(expval))
             resid = (endog[i] - expval) / stdev
 
-            j1, j2 = np.tril_indices(len(expval))
+            j1, j2 = np.tril_indices(len(expval), -1)
             dx = np.abs(self.time[i][j1] - self.time[i][j2])
             ii = np.flatnonzero(dx <= self.max_lag)
             j1 = j1[ii]
             j2 = j2[ii]
             dx = dx[ii]
 
-            vs = np.bincount(dx, weights=resid[
-                             j1] * resid[j2], minlength=self.max_lag + 1)
+            vs = np.bincount(dx, weights=resid[j1] * resid[j2],
+                             minlength=self.max_lag + 1)
             vd = np.bincount(dx, minlength=self.max_lag + 1)
 
+            resid_ssq += np.sum(resid**2)
+            resid_ssq_n += len(resid)
+
             ii = np.flatnonzero(vd > 0)
             if len(ii) > 0:
                 dn[ii] += 1
                 dep_params[ii] += vs[ii] / vd[ii]
 
         i0 = np.flatnonzero(dn > 0)
         dep_params[i0] /= dn[i0]
-        self.dep_params = dep_params[1:] / dep_params[0]
+        resid_msq = resid_ssq / resid_ssq_n
+        dep_params /= resid_msq
+        self.dep_params = dep_params
 
     def covariance_matrix(self, endog_expval, index):
 
         if self.grid:
             return self.covariance_matrix_grid(endog_expval, index)
 
-        j1, j2 = np.tril_indices(len(endog_expval))
+        j1, j2 = np.tril_indices(len(endog_expval), -1)
         dx = np.abs(self.time[index][j1] - self.time[index][j2])
-        ii = np.flatnonzero((0 < dx) & (dx <= self.max_lag))
+        ii = np.flatnonzero(dx <= self.max_lag)
         j1 = j1[ii]
         j2 = j2[ii]
         dx = dx[ii]
 
         cmat = np.eye(len(endog_expval))
-        cmat[j1, j2] = self.dep_params[dx - 1]
-        cmat[j2, j1] = self.dep_params[dx - 1]
+        cmat[j1, j2] = self.dep_params[dx]
+        cmat[j2, j1] = self.dep_params[dx]
+
         return cmat, True
 
     def covariance_matrix_grid(self, endog_expval, index):
 
         from scipy.linalg import toeplitz
         r = np.zeros(len(endog_expval))
         r[0] = 1
-        r[1:self.max_lag + 1] = self.dep_params
+        r[1:self.max_lag + 1] = self.dep_params[1:]
         return toeplitz(r), True
 
     def covariance_matrix_solve(self, expval, index, stdev, rhs):
@@ -629,7 +641,7 @@ def covariance_matrix_solve(self, expval, index, stdev, rhs):
 
         from statsmodels.tools.linalg import stationary_solve
         r = np.zeros(len(expval))
-        r[0:self.max_lag] = self.dep_params
+        r[0:self.max_lag] = self.dep_params[1:]
         return [stationary_solve(r, x) for x in rhs]
 
     update.__doc__ = CovStruct.update.__doc__
@@ -638,8 +650,8 @@ def covariance_matrix_solve(self, expval, index, stdev, rhs):
 
     def summary(self):
 
-        return ("Stationary dependence parameters\n",
-                self.dep_params)
+        lag = np.arange(self.max_lag + 1)
+        return pd.DataFrame({"Lag": lag, "Cov": self.dep_params})
 
 
 class Autoregressive(CovStruct):