Fix two issues

econml/dml.py

@@ -652,7 +652,7 @@ def fit(self, Y, T, X=None, W=None, sample_weight=None, inference=None):
         inference: string, `Inference` instance, or None
             Method for performing inference.  This estimator supports 'bootstrap'
             (or an instance of :class:`.BootstrapInference`) and 'debiasedlasso'
-            (or an instance of :class:`.LinearCateInference`)
+            (or an instance of :class:`.LinearModelFinalInference`)
 
         Returns
         -------

econml/inference.py

@@ -92,7 +92,7 @@ def const_marginal_effect_interval(self, X, *, alpha=0.1):
         if X is None:
             X = np.ones((1, 1))
         elif self.featurizer is not None:
-            X = self.featurizer.fit_transform(X)
+            X = self.featurizer.transform(X)
         X, T = broadcast_unit_treatments(X, self.d_t)
         preds = self._predict_interval(cross_product(X, T), alpha=alpha)
         return tuple(reshape_treatmentwise_effects(pred, self._d_t, self._d_y)
@@ -102,7 +102,7 @@ def const_marginal_effect_inference(self, X):
         if X is None:
             X = np.ones((1, 1))
         elif self.featurizer is not None:
-            X = self.featurizer.fit_transform(X)
+            X = self.featurizer.transform(X)
         X, T = broadcast_unit_treatments(X, self.d_t)
         pred = reshape_treatmentwise_effects(self._predict(cross_product(X, T)), self._d_t, self._d_y)
         if not hasattr(self.model_final, 'prediction_stderr'):
@@ -356,13 +356,13 @@ def fit(self, estimator, *args, **kwargs):
 
     def const_marginal_effect_interval(self, X, *, alpha=0.1):
         if (X is not None) and (self.featurizer is not None):
-            X = self.featurizer.fit_transform(X)
+            X = self.featurizer.transform(X)
         preds = np.array([mdl.predict_interval(X, alpha=alpha) for mdl in self.fitted_models_final])
         return tuple(np.moveaxis(preds, [0, 1], [-1, 0]))  # send treatment to the end, pull bounds to the front
 
     def const_marginal_effect_inference(self, X):
         if (X is not None) and (self.featurizer is not None):
-            X = self.featurizer.fit_transform(X)
+            X = self.featurizer.transform(X)
         pred = np.array([mdl.predict(X) for mdl in self.fitted_models_final])
         if not hasattr(self.fitted_models_final[0], 'prediction_stderr'):
             raise AttributeError("Final model doesn't support prediction standard eror, "

econml/tests/test_dml.py

@@ -693,6 +693,30 @@ def test_can_custom_splitter(self):
         dml.fit(np.array([1, 2, 3, 1, 2, 3]), np.array([1, 2, 3, 1, 2, 3]), np.ones((6, 1)))
         dml.score(np.array([1, 2, 3, 1, 2, 3]), np.array([1, 2, 3, 1, 2, 3]), np.ones((6, 1)))
 
+    def test_can_use_featurizer(self):
+        "Test that we can use a featurizer, and that fit is only called during training"
+        dml = LinearDMLCateEstimator(LinearRegression(), LinearRegression(),
+                                     fit_cate_intercept=False, featurizer=OneHotEncoder(n_values='auto', sparse=False))
+
+        T = np.tile([1, 2, 3], 6)
+        Y = np.array([1, 2, 3, 1, 2, 3])
+        Y = np.concatenate([Y, 0 * Y, -Y])
+        X = np.repeat([[7, 8, 9]], 6, axis=1).T
+
+        dml.fit(Y, T, X=X, inference='statsmodels')
+
+        # because there is one fewer unique element in the test set, fit_transform would return the wrong number of fts
+        X_test = np.array([[7, 8]]).T
+
+        np.testing.assert_equal(dml.effect(X_test)[::-1], dml.effect(X_test[::-1]))
+        eff_int = np.array(dml.effect_interval(X_test))
+        eff_int_rev = np.array(dml.effect_interval(X_test[::-1]))
+        np.testing.assert_equal(eff_int[:, ::-1], eff_int_rev)
+
+        eff_int = np.array(dml.const_marginal_effect_interval(X_test))
+        eff_int_rev = np.array(dml.const_marginal_effect_interval(X_test[::-1]))
+        np.testing.assert_equal(eff_int[:, ::-1], eff_int_rev)
+
     def test_can_use_statsmodel_inference(self):
         """Test that we can use statsmodels to generate confidence intervals"""
         dml = LinearDMLCateEstimator(LinearRegression(), LogisticRegression(C=1000),

econml/two_stage_least_squares.py

@@ -294,9 +294,9 @@ def effect(self, X=None, T0=0, T1=1):
         assert shape(T0)[0] == shape(X)[0]
 
         W = np.zeros((shape(T0)[0], self._d_w))  # can set arbitrarily since values will cancel
-        ft_X = self._x_featurizer.fit_transform(X)
-        ft_T0 = self._t_featurizer.fit_transform(T0)
-        ft_T1 = self._t_featurizer.fit_transform(T1)
+        ft_X = self._x_featurizer.transform(X)
+        ft_T0 = self._t_featurizer.transform(T0)
+        ft_T1 = self._t_featurizer.transform(T1)
         Y0 = self._model_Y.predict(_add_ones(np.hstack([W, cross_product(ft_T0, ft_X)])))
         Y1 = self._model_Y.predict(_add_ones(np.hstack([W, cross_product(ft_T1, ft_X)])))
         return Y1 - Y0
@@ -327,9 +327,9 @@ def marginal_effect(self, T, X=None):
             X = np.empty((shape(T)[0], 0))
         assert shape(T)[0] == shape(X)[0]
 
-        ft_X = self._x_featurizer.fit_transform(X)
+        ft_X = self._x_featurizer.transform(X)
         n = shape(T)[0]
-        dT = self._dt_featurizer.fit_transform(T)
+        dT = self._dt_featurizer.transform(T)
         W = np.zeros((size(T), self._d_w))
         # dT should be an n×dₜ×fₜ array (but if T was a vector, or if there is only one feature,
         # dT may be only 2-dimensional)