ENH: make method signature more consistent with new statsmodels behavior. Uses dot product directly so pandas users aren't affected by statsmodels API change

Author: Chang She

pandas/stats/ols.py

@@ -9,7 +9,7 @@
 
 import numpy as np
 
-from pandas.core.api import DataFrame, Series
+from pandas.core.api import DataFrame, Series, isnull
 from pandas.core.common import _ensure_float64
 from pandas.core.index import MultiIndex
 from pandas.core.panel import Panel
@@ -381,12 +381,13 @@ def y_predict(self):
         For in-sample, this is same as y_fitted."""
         return self.y_fitted
 
-    def predict(self, new_y_values, fill_value=None, fill_method=None,
-                axis=0):
+    def predict(self, beta=None, x=None, fill_value=None,
+                fill_method=None, axis=0):
         """
         Parameters
         ----------
-        new_y_values : Series or DataFrame
+        beta : Series
+        x : Series or DataFrame
         fill_value : scalar or dict, default None
         fill_method : {'backfill', 'bfill', 'pad', 'ffill', None}, default None
         axis : {0, 1}, default 0
@@ -403,20 +404,34 @@ def predict(self, new_y_values, fill_value=None, fill_method=None,
         -------
         Series of predicted values
         """
-        orig_y = new_y_values
-        if fill_value is None and fill_method is None:
-            new_y_values = new_y_values.dropna(how='any')
+        if beta is None and x is None:
+            return self.y_predict
+
+        if beta is None:
+            beta = self.beta
         else:
-            new_y_values = new_y_values.fillna(value=fill_value,
-                                               method=fill_method, axis=axis)
-        if isinstance(new_y_values, Series):
-            new_y_values = DataFrame({'x' : new_y_values})
-        if self._intercept:
-            new_y_values['intercept'] = 1.
+            beta = beta.reindex(self.beta.index)
+            if isnull(beta).any():
+                raise ValueError('Must supply betas for same variables')
+
+        if x is None:
+            x = self._x
+            orig_x = x
+        else:
+            orig_x = x
+            if fill_value is None and fill_method is None:
+                x = x.dropna(how='any')
+            else:
+                x = x.fillna(value=fill_value, method=fill_method, axis=axis)
+            if isinstance(x, Series):
+                x = DataFrame({'x' : x})
+            if self._intercept:
+                x['intercept'] = 1.
+
+            x = x.reindex(columns=self._x.columns)
 
-        new_y_values = new_y_values.reindex(columns=self._x.columns)
-        rs = self.sm_ols.model.predict(new_y_values.values)
-        return Series(rs, new_y_values.index).reindex(orig_y.index)
+        rs = x.values.dot(beta.values)
+        return Series(rs, x.index).reindex(orig_x.index)
 
     RESULT_FIELDS = ['r2', 'r2_adj', 'df', 'df_model', 'df_resid', 'rmse',
                      'f_stat', 'beta', 'std_err', 't_stat', 'p_value', 'nobs']

pandas/stats/tests/test_ols.py

@@ -291,17 +291,29 @@ def test_predict(self):
         y = tm.makeTimeSeries()
         x = tm.makeTimeDataFrame()
         model1 = ols(y=y, x=x)
-        assert_series_equal(model1.predict(x), model1.y_predict)
+        assert_series_equal(model1.predict(), model1.y_predict)
+        assert_series_equal(model1.predict(x=x), model1.y_predict)
+        assert_series_equal(model1.predict(beta=model1.beta), model1.y_predict)
+
+        exog = x.copy()
+        exog['intercept'] = 1.
+        rs = Series(exog.values.dot(model1.beta.values), x.index)
+        assert_series_equal(model1.y_predict, rs)
+
         x2 = x.reindex(columns=x.columns[::-1])
-        assert_series_equal(model1.predict(x2), model1.y_predict)
+        assert_series_equal(model1.predict(x=x2), model1.y_predict)
 
         x3 = x2 + 10
-        pred3 = model1.predict(x3)
+        pred3 = model1.predict(x=x3)
         x3['intercept'] = 1.
         x3 = x3.reindex(columns = model1.beta.index)
         expected = Series(x3.values.dot(model1.beta.values), x3.index)
         assert_series_equal(expected, pred3)
 
+        beta = Series(0., model1.beta.index)
+        pred4 = model1.predict(beta=beta)
+        assert_series_equal(Series(0., pred4.index), pred4)
+
     def test_longpanel_series_combo(self):
         wp = tm.makePanel()
         lp = wp.to_frame()