Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions linearmodels/iv/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down
38 changes: 37 additions & 1 deletion linearmodels/tests/iv/test_formulas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down Expand Up @@ -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)
35 changes: 34 additions & 1 deletion linearmodels/tests/iv/test_results.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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"]
Loading