Prediction error for non-fe model

justfile

@@ -1,5 +1,5 @@
 # uncomment for windows
-set shell := ["powershell.exe", "-c"]
+# set shell := ["powershell.exe", "-c"]
 
 # docs: justfile docs: https://just.systems/man/en/
 default:

pyfixest/did/did2s.py

@@ -216,7 +216,7 @@ def _did2s_estimate(
     fit1.fixef()
 
     # demean data
-    Y_hat = fit1.predict(newdata=data)
+    Y_hat = fit1.predict(newdata=data).yhat
     _first_u = data[f"{yname}"].to_numpy().flatten() - Y_hat
     data[f"{yname}_hat"] = _first_u
 

pyfixest/estimation/feols_.py

@@ -1426,7 +1426,8 @@
         atol: float = 1e-6,
         btol: float = 1e-6,
         type: str = "link",
-    ) -> np.ndarray:
+        compute_stdp: bool = False
+    ) -> pd.DataFrame:
         """
         Predict values of the model on new data.
 
@@ -1451,23 +1452,29 @@
             Another stopping tolerance for scipy.sparse.linalg.lsqr().
             See https://docs.scipy.org/doc/
                 scipy/reference/generated/scipy.sparse.linalg.lsqr.html
-        link:
+        type:
             The type of prediction to be made. Can be either 'link' or 'response'.
              Defaults to 'link'. 'link' and 'response' lead
             to identical results for linear models.
+        compute_stdp: boolean
+            Whether to compute standard error of the predictions
 
         Returns
         -------
-        y_hat : np.ndarray
-            A flat np.array with predicted values of the regression model.
+        pred_results : pd.DataFrame
+            Dataframe with columns "y_hat" and optionally "stdp" (if compute_stdp is True)
         """
         if self._is_iv:
             raise NotImplementedError(
                 "The predict() method is currently not supported for IV models."
             )
 
+        prediction_df = pd.DataFrame()
         if newdata is None:
-            return self._Y_untransformed.to_numpy().flatten() - self.resid()
+            prediction_df["yhat"] = self._Y_untransformed.to_numpy().flatten() - self._u_hat.flatten()
+            if compute_stdp:
+                prediction_df["stdp"] = self.get_newdata_stdp(self._X)
+            return prediction_df
 
         newdata = _polars_to_pandas(newdata).reset_index(drop=False)
 
@@ -1501,7 +1508,64 @@
             _fixef_mat = _apply_fixef_numpy(df_fe.values, fixef_dicts)
             y_hat += np.sum(_fixef_mat, axis=1)
 
-        return y_hat.flatten()
+        prediction_df["yhat"] = y_hat.flatten()
+        if compute_stdp:
+            stdp_df = np.full(newdata.shape[0], np.nan)
+            stdp_df[X_index] = self.get_newdata_stdp(X)
+            prediction_df["stdp"] = stdp_df
+
+        return prediction_df
+
+
+    def get_newdata_stdp(self, X):
+        """
+        Get standard error of predictions for new data.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            Covariates for new data points.
+
+        Returns
+        -------
+        list
+            Standard errors for each prediction
+        """
+        # for now only compute prediction error if model has no fixed effects
+        # TODO: implement for fixed effects
+        if not self._has_fixef:
+            if not self._X_is_empty:
+                return list(map(self.get_single_row_stdp, X))
+            else:
+                warnings.warn(
+                    """
+                    Standard error of the prediction cannot be computed if X is empty.
+                    Prediction dataframe stdp column will be None.
+                    """
+                )
+        else:
+            warnings.warn(
+                """
+                Standard error of the prediction is not implemented for fixed effects models.
+                Prediction dataframe stdp column will be None.
+                """
+            )
+
+    def get_single_row_stdp(self, row):
+        """
+        Get standard error of predictions for a single row.
+
+        Parameters
+        ----------
+        row : np.ndarray
+            Single row of new covariate data
+
+        Returns
+        -------
+        np.ndarray
+            Standard error of prediction for single row
+        """
+        return np.linalg.multi_dot([row, self._vcov, np.transpose(row)])
 
     def get_nobs(self):
         """

pyfixest/estimation/fepois_.py

@@ -299,7 +299,8 @@
         atol: float = 1e-6,
         btol: float = 1e-6,
         type: str = "link",
-    ) -> np.ndarray:
+        compute_stdp: bool = False
+    ) -> pd.DataFrame:
         """
         Return predicted values from regression model.
 
@@ -334,13 +335,13 @@
         btol : Float, default 1e-6
             Another stopping tolerance for scipy.sparse.linalg.lsqr().
             See https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.lsqr.html
-
-
+        compute_stdp: boolean
+            Ignored, included to conform with method signature of feols.predict()
 
         Returns
         -------
-        np.ndarray
-            A flat array with the predicted values of the regression model.
+        pred_results : pd.DataFrame
+            Dataframe with columns "y_hat" with predicted values from regression model
         """
         _Xbeta = self._Xbeta.flatten()
         _has_fixef = self._has_fixef
@@ -355,10 +356,13 @@
             )
 
         # y_hat = super().predict(newdata=newdata, type=type, atol=atol, btol=btol)
+        prediction_df = pd.DataFrame()
         if type == "link":
-            return np.exp(_Xbeta)
+            prediction_df["yhat"] = np.exp(_Xbeta)
+            return prediction_df
         elif type == "response":
-            return _Xbeta
+            prediction_df["yhat"] = _Xbeta
+            return prediction_df
         else:
             raise ValueError("type must be one of 'response' or 'link'.")
 

pyproject.toml

@@ -37,6 +37,7 @@ wildboottest = ">=0.2.0"
 pre-commit = "^3.6.0"
 doubleml = "^0.7.1"
 marginaleffects = "^0.0.10"
+statsmodels = "^0.14.1"
 
 [tool.poetry.group.docs.dependencies]
 causaldata = "*"

tests/test_predict_resid_fixef.py

@@ -37,8 +37,8 @@ def test_internally(data):
     # predict via feols, without fixed effect
     fit = feols(fml="Y~csw(X1, X2)", data=data, vcov="iid")
     mod = fit.fetch_model(0)
-    original_prediction = mod.predict()
-    updated_prediction = mod.predict(newdata=mod._data)
+    original_prediction = mod.predict().yhat
+    updated_prediction = mod.predict(newdata=mod._data).yhat
     np.allclose(original_prediction, updated_prediction)
     assert mod._data.shape[0] == original_prediction.shape[0]
     assert mod._data.shape[0] == updated_prediction.shape[0]
@@ -47,15 +47,15 @@ def test_internally(data):
     fit = feols(fml="Y~csw(X1, X2) | f1", data=data, vcov="iid")
     mod = fit.fetch_model(0)
     mod.fixef()
-    original_prediction = mod.predict()
-    updated_prediction = mod.predict(newdata=mod._data)
+    original_prediction = mod.predict().yhat
+    updated_prediction = mod.predict(newdata=mod._data).yhat
     np.allclose(original_prediction, updated_prediction)
     assert mod._data.shape[0] == original_prediction.shape[0]
     assert mod._data.shape[0] == updated_prediction.shape[0]
 
     # now expect error with updated predicted being a subset of data
     with pytest.raises(ValueError):
-        updated_prediction = mod.predict(newdata=data.iloc[0:100, :])
+        updated_prediction = mod.predict(newdata=data.iloc[0:100, :]).yhat
         np.allclose(original_prediction, updated_prediction)
 
     # fepois NotImplementedError(s)
@@ -123,18 +123,18 @@ def test_vs_fixest(data, fml):
     #    raise ValueError("sumFE for Poisson are not equal")
 
     # test predict for OLS
-    if not np.allclose(feols_mod.predict(), r_fixest_ols.rx2("fitted.values")):
+    if not np.allclose(feols_mod.predict().yhat, r_fixest_ols.rx2("fitted.values")):
         raise ValueError("Predictions for OLS are not equal")
 
-    if not np.allclose(len(feols_mod.predict()), len(stats.predict(r_fixest_ols))):
+    if not np.allclose(len(feols_mod.predict().yhat), len(stats.predict(r_fixest_ols))):
         raise ValueError("Predictions for OLS are not the same length")
     # test predict for Poisson
-    # if not np.allclose(fepois_mod.predict(), r_fixest_pois.rx2("fitted.values")):
+    # if not np.allclose(fepois_mod.predict().yhat, r_fixest_pois.rx2("fitted.values")):
     #    raise ValueError("Predictions for Poisson are not equal")
 
     # test on new data - OLS.
     if not np.allclose(
-        feols_mod.predict(newdata=data2), stats.predict(r_fixest_ols, newdata=data2)
+        feols_mod.predict(newdata=data2).yhat, stats.predict(r_fixest_ols, newdata=data2)
     ):
         raise ValueError("Predictions for OLS are not equal")
 
@@ -171,7 +171,7 @@ def test_predict_nas():
     # test 1
     fml = "Y ~ X1 + X2 | f1"
     fit = feols(fml, data=data)
-    res = fit.predict(newdata=data)
+    res = fit.predict(newdata=data).yhat
     fit_r = fixest.feols(ro.Formula(fml), data=data)
     res_r = stats.predict(fit_r, newdata=data)
     np.testing.assert_allclose(res, res_r, atol=1e-05, rtol=1e-05)
@@ -184,15 +184,15 @@ def test_predict_nas():
 
     fml = "Y ~ X1 + X2 | f1"
     fit = feols(fml, data=data)
-    res = fit.predict(newdata=newdata)
+    res = fit.predict(newdata=newdata).yhat
     fit_r = fixest.feols(ro.Formula(fml), data=data)
     res_r = stats.predict(fit_r, newdata=newdata)
     np.testing.assert_allclose(res, res_r, atol=1e-05, rtol=1e-05)
     assert newdata.shape[0] == len(res)
     assert len(res) == len(res_r)
 
     newdata.loc[198, "Y"] = np.nan
-    res = fit.predict(newdata=newdata)
+    res = fit.predict(newdata=newdata).yhat
     res_r = stats.predict(fit_r, newdata=newdata)
     np.testing.assert_allclose(res, res_r, atol=1e-05, rtol=1e-05)
     assert newdata.shape[0] == len(res)
@@ -201,7 +201,7 @@ def test_predict_nas():
     # test 3
     fml = "Y ~ X1 + X2 + f1| f1 "
     fit = feols(fml, data=data)
-    res = fit.predict(newdata=data)
+    res = fit.predict(newdata=data).yhat
     fit_r = fixest.feols(ro.Formula(fml), data=data)
     res_r = stats.predict(fit_r, newdata=data)
     np.testing.assert_allclose(res, res_r, atol=1e-05, rtol=1e-05)
@@ -229,7 +229,7 @@ def test_new_fixef_level(data, fml):
         se="hetero",
     )
 
-    updated_prediction_py = feols_mod.predict(newdata=data2)
+    updated_prediction_py = feols_mod.predict(newdata=data2).yhat
     updated_prediction_r = stats.predict(r_fixest_ols, newdata=data2)
 
     if not np.allclose(updated_prediction_py, updated_prediction_r):
@@ -249,7 +249,7 @@ def test_categorical_covariate_predict():
     df_sub = df.query("x == 1 or x == 2 or x == 3").copy()
 
     py_fit = feols("y ~ C(x, contr.treatment(base=1))", df)
-    py_predict = py_fit.predict(df_sub)
+    py_predict = py_fit.predict(df_sub).yhat
 
     r_predict = np.array(
         [

tests/test_stdp.py

@@ -0,0 +1,22 @@
+import numpy as np
+import statsmodels.api as sm
+
+from pyfixest.estimation.estimation import feols
+from pyfixest.utils.utils import get_data
+
+
+def test_stdp():
+    """Compare the standard error of the prediction to statsmodels.get_prediction()."""
+    data = get_data().dropna()
+
+    fit = feols("Y ~ X1 + X2", data=data)
+    stdp = fit.predict(compute_stdp=True).stdp
+
+    X = data[["X1","X2"]]
+    Y = data["Y"]
+    X = sm.add_constant(X)
+    model = sm.OLS(Y, X).fit()
+    predictions = model.get_prediction()
+    stdp_sm = predictions.var_pred_mean
+
+    np.testing.assert_allclose(stdp, stdp_sm)