Merge pull request #677 from bashtage/fix-predict-none

BUG: Fix predict when exog or endog is None

Author: bashtage

linearmodels/iv/model.py

@@ -291,15 +291,24 @@ def predict(
                 "Predictions can only be constructed using one "
                 "of exog/endog or data, but not both."
             )
-        if exog is not None or endog is not None:
+        if exog is not None:
             exog = IVData(exog).pandas
+        if endog is not None:
             endog = IVData(endog).pandas
-        elif data is not None:
+        if data is not None:
             parser = IVFormulaParser(self.formula, data, eval_env=eval_env)
             exog = parser.exog
             endog = parser.endog
-        else:
-            raise ValueError("exog and endog or data must be provided.")
+        if all(a is None for a in (exog, endog, data)):
+            raise ValueError("At least one of exog, endog, or data must be provided.")
+        if exog is None:
+            assert endog is not None
+            exog = IVData(exog, nobs=endog.shape[0]).pandas
+            exog.index = endog.index
+        if endog is None:
+            assert exog is not None
+            endog = IVData(endog, nobs=exog.shape[0]).pandas
+            endog.index = exog.index
         assert exog is not None
         assert endog is not None
         if exog.shape[0] != endog.shape[0]:

linearmodels/tests/iv/test_formulas.py

@@ -210,7 +210,9 @@ def test_predict_formula(data, model_and_func, formula):
     assert_frame_equal(pred, pred2)
     assert_allclose(res.fitted_values, pred)
 
-    with pytest.raises(ValueError, match=r"exog and endog or data must be provided"):
+    with pytest.raises(
+        ValueError, match=r"At least one of exog, endog, or data must be provided"
+    ):
         mod.predict(res.params)
 
 
@@ -404,3 +406,37 @@ def test_formula_categorical_equiv(data, model_and_func, dtype):
         "x2",
         "x3",
     ]
+
+
+def test_predict_no_rhs(data, model_and_func):
+    model, _ = model_and_func
+    mod = model.from_formula("y ~", data)
+    res = mod.fit()
+    pred0 = res.predict()
+    pred1 = res.predict(data=data)
+    pred1.columns = pred0.columns
+    assert_frame_equal(pred0, pred1)
+
+
+@pytest.mark.parametrize(
+    "fmla",
+    [
+        "y ~ 1",
+        "y ~ 1 + x3",
+        "y ~ [x1 + x2 ~ 1 + z1 + z2 + z3]",
+        "y ~ 1 + [x1 + x2 ~ z1 + z2 + z3]",
+    ],
+)
+def test_formula_single(data, model_and_func, fmla):
+    model, func = model_and_func
+    res = model.from_formula(fmla, data).fit()
+    pred0 = res.predict()
+    pred1 = res.predict(data=data)
+
+    mod2 = func(fmla, data)
+    res2 = mod2.fit()
+    pred2 = res2.predict(data=data)
+    pred1.columns = pred0.columns
+    pred2.columns = pred0.columns
+    assert_frame_equal(pred1, pred2)
+    assert_frame_equal(pred0, pred1)

linearmodels/tests/iv/test_results.py

@@ -1,7 +1,7 @@
 from numpy import asarray
 from numpy.testing import assert_allclose
 from pandas import DataFrame
-from pandas.testing import assert_series_equal
+from pandas.testing import assert_frame_equal, assert_series_equal
 import pytest
 
 from linearmodels.iv.data import IVData
@@ -96,3 +96,36 @@ def test_predict_no_selection(data, model):
     res = mod.fit()
     with pytest.raises(ValueError, match=r"At least one output must be selected"):
         res.predict(fitted=False, idiosyncratic=False, missing=True)
+
+
+@pytest.mark.parametrize("include_vars", ["exog", "both", "endog"])
+def test_fitted_predict_combinations(data, model, include_vars):
+    args = (data.dep,)
+    if include_vars in ("exog", "both"):
+        args += (data.exog,)
+    else:
+        args += (None,)
+    if include_vars in ("endog", "both"):
+        args += (data.endog, data.instr)
+    else:
+        args += (None, None)
+
+    mod = model(*args)
+    res = mod.fit()
+    assert_series_equal(res.idiosyncratic, res.resids)
+    y = mod.dependent.pandas
+    expected = asarray(y) - asarray(res.resids)[:, None]
+    expected = DataFrame(expected, y.index, ["fitted_values"])
+    assert_frame_similar(expected, res.fitted_values)
+    assert_allclose(expected, res.fitted_values)
+    pred = res.predict()
+    pred2 = res.predict(exog=args[1], endog=args[2])
+    pred2.columns = pred.columns
+    assert_frame_equal(pred, pred2)
+    nobs = res.resids.shape[0]
+    assert isinstance(pred, DataFrame)
+    assert pred.shape == (nobs, 1)
+    pred = res.predict(idiosyncratic=True, missing=True)
+    nobs = IVData(data.dep).pandas.shape[0]
+    assert pred.shape == (nobs, 2)
+    assert list(pred.columns) == ["fitted_values", "residual"]