fix test error

Author: heimengqi

README.md

@@ -290,8 +290,8 @@ from econml.iv.dr import LinearIntentToTreatDRIV
 from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
 from sklearn.linear_model import LinearRegression
 
-est = LinearIntentToTreatDRIV(model_Y_X=GradientBoostingRegressor(),
-                              model_T_XZ=GradientBoostingClassifier(),
+est = LinearIntentToTreatDRIV(model_y_xw=GradientBoostingRegressor(),
+                              model_t_xwz=GradientBoostingClassifier(),
                               flexible_model_effect=GradientBoostingRegressor())
 est.fit(Y, T, Z=Z, X=X) # OLS inference by default
 treatment_effects = est.effect(X_test)

econml/iv/dr/_dr.py

@@ -29,7 +29,7 @@
 from ...sklearn_extensions.linear_model import StatsModelsLinearRegression, DebiasedLasso, WeightedLassoCVWrapper
 from ...sklearn_extensions.model_selection import WeightedStratifiedKFold
 from ...utilities import (_deprecate_positional, add_intercept, filter_none_kwargs,
-                          inverse_onehot, get_feature_names_or_default, check_high_dimensional)
+                          inverse_onehot, get_feature_names_or_default, check_high_dimensional, check_input_arrays)
 from ...grf import RegressionForest
 from ...dml.dml import _FirstStageWrapper, _FinalWrapper
 from ..._shap import _shap_explain_model_cate
@@ -176,7 +176,6 @@ def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None)
         Y_pred = Y_pred.reshape(Y.shape)
         T_pred = T_pred.reshape(T.shape)
         TZ_pred = TZ_pred.reshape(T.shape)
-        prel_theta = prel_theta.reshape(Y.shape)
 
         Y_res = Y - Y_pred
         T_res = T - T_pred
@@ -196,8 +195,7 @@ def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None)
             cov = TZ_pred - T_pred * Z_pred
 
         # check nuisances outcome shape
-        assert prel_theta.ndim == 1, "Nuisance outcome should be vector!"
-        assert Y_res.ndim == 1, "Nuisance outcome should be vector!"
+        # Y_res could be a vector or 1-dimensional 2d-array
         assert T_res.ndim == 1, "Nuisance outcome should be vector!"
         assert Z_res.ndim == 1, "Nuisance outcome should be vector!"
         assert cov.ndim == 1, "Nuisance outcome should be vector!"
@@ -340,21 +338,19 @@ def _gen_ortho_learner_model_final(self):
                                    self.cov_clip, self.opt_reweighted)
 
     def _check_inputs(self, Y, T, Z, X, W):
-        if len(Y.shape) > 1 and Y.shape[1] > 1:
+        Y1, T1, Z1, = check_input_arrays(Y, T, Z)
+        if len(Y1.shape) > 1 and Y1.shape[1] > 1:
             raise AssertionError("DRIV only supports single dimensional outcome")
-        if len(T.shape) > 1 and T.shape[1] > 1:
+        if len(T1.shape) > 1 and T1.shape[1] > 1:
             if self.discrete_treatment:
                 raise AttributeError("DRIV only supports binary treatments")
             else:
                 raise AttributeError("DRIV only supports single-dimensional continuous treatments")
-        if len(Z.shape) > 1 and Z.shape[1] > 1:
+        if len(Z1.shape) > 1 and Z1.shape[1] > 1:
             if self.discrete_instrument:
                 raise AttributeError("DRIV only supports binary instruments")
             else:
                 raise AttributeError("DRIV only supports single-dimensional continuous instruments")
-        Z = Z.ravel()
-        T = T.ravel()
-        Y = Y.ravel()
         return Y, T, Z, X, W
 
     @_deprecate_positional("X and W should be passed by keyword only. In a future release "
@@ -1746,19 +1742,14 @@ def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None)
         Z_pred = Z_pred.reshape(Z.shape)
         T_pred_one = T_pred_one.reshape(T.shape)
         T_pred_zero = T_pred_zero.reshape(T.shape)
-        prel_theta = prel_theta.reshape(Y.shape)
 
+        # T_res, Z_res, beta expect shape to be (n,1)
         beta = Z_pred * (1 - Z_pred) * (T_pred_one - T_pred_zero)
         T_pred = T_pred_one * Z_pred + T_pred_zero * (1 - Z_pred)
         Y_res = Y - Y_pred
         T_res = T - T_pred
         Z_res = Z - Z_pred
 
-        # check nuisances outcome shape
-        # T_res, Z_res, beta expect shape to be (n,1)
-        assert prel_theta.ndim == 1, "Nuisance outcome should be vector!"
-        assert Y_res.ndim == 1, "Nuisance outcome should be vector!"
-
         return prel_theta, Y_res, T_res, Z_res, beta
 
 

econml/ortho_iv.py

@@ -6,15 +6,15 @@
 from .utilities import deprecated
 
 
-@deprecated("The econml.ortho_iv.DMLATEIV class has been moved to econml.iv.dml.DMLATEIV; "
+@deprecated("The econml.ortho_iv.DMLATEIV class has been moved to econml.iv.dml.OrthoIV; "
             "an upcoming release will remove support for the old name")
-class DMLATEIV(dmliv.DMLATEIV):
+class DMLATEIV(dmliv.OrthoIV):
     pass
 
 
-@deprecated("The econml.ortho_iv.ProjectedDMLATEIV class has been moved to econml.iv.dml.ProjectedDMLATEIV; "
+@deprecated("The econml.ortho_iv.ProjectedDMLATEIV class has been moved to econml.iv.dml.OrthoIV; "
             "an upcoming release will remove support for the old name")
-class ProjectedDMLATEIV(dmliv.ProjectedDMLATEIV):
+class ProjectedDMLATEIV(dmliv.OrthoIV):
     pass
 
 

econml/tests/test_integration.py

@@ -16,7 +16,7 @@
 from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
 from sklearn.linear_model import LinearRegression, MultiTaskLasso, LassoCV
 from sklearn.preprocessing import PolynomialFeatures, FunctionTransformer
-from econml.ortho_iv import LinearIntentToTreatDRIV
+from econml.iv.dr import LinearIntentToTreatDRIV
 from econml.deepiv import DeepIVEstimator