Refitting model_final and nuisance averaging

econml/_cate_estimator.py

@@ -529,8 +529,11 @@ class LinearModelFinalCateEstimatorMixin(BaseCateEstimator):
     """
     Base class for models where the final stage is a linear model.
 
-    Subclasses must expose a ``model_final`` attribute containing the model's
-    final stage model.
+    Such an estimator must implement a :attr:`model_final_` attribute that points
+    to the fitted final :class:`.StatsModelsLinearRegression` object that
+    represents the fitted CATE model. Also must implement :attr:`featurizer_` that points
+    to the fitted featurizer and :attr:`bias_part_of_coef` that designates
+    if the intercept is the first element of the :attr:`model_final_` coefficient.
 
     Attributes
     ----------
@@ -544,7 +547,9 @@ def _get_inference_options(self):
         options.update(auto=LinearModelFinalInference)
         return options
 
-    bias_part_of_coef = False
+    @property
+    def bias_part_of_coef(self):
+        return False
 
     @property
     def coef_(self):
@@ -561,9 +566,9 @@ def coef_(self):
             a vector and not a 2D array. For binary treatment the n_t dimension is
             also omitted.
         """
-        return parse_final_model_params(self.model_final.coef_, self.model_final.intercept_,
+        return parse_final_model_params(self.model_final_.coef_, self.model_final_.intercept_,
                                         self._d_y, self._d_t, self._d_t_in, self.bias_part_of_coef,
-                                        self.fit_cate_intercept)[0]
+                                        self.fit_cate_intercept_)[0]
 
     @property
     def intercept_(self):
@@ -578,11 +583,11 @@ def intercept_(self):
             a vector and not a 2D array. For binary treatment the n_t dimension is
             also omitted.
         """
-        if not self.fit_cate_intercept:
+        if not self.fit_cate_intercept_:
             raise AttributeError("No intercept was fitted!")
-        return parse_final_model_params(self.model_final.coef_, self.model_final.intercept_,
+        return parse_final_model_params(self.model_final_.coef_, self.model_final_.intercept_,
                                         self._d_y, self._d_t, self._d_t_in, self.bias_part_of_coef,
-                                        self.fit_cate_intercept)[1]
+                                        self.fit_cate_intercept_)[1]
 
     @BaseCateEstimator._defer_to_inference
     def coef__interval(self, *, alpha=0.1):
@@ -718,11 +723,11 @@ def summary(self, alpha=0.1, value=0, decimals=3, feature_names=None, treatment_
             return smry
 
     def shap_values(self, X, *, feature_names=None, treatment_names=None, output_names=None, background_samples=100):
-        if hasattr(self, "featurizer") and self.featurizer is not None:
-            X = self.featurizer.transform(X)
+        if hasattr(self, "featurizer_") and self.featurizer_ is not None:
+            X = self.featurizer_.transform(X)
         feature_names = self.cate_feature_names(feature_names)
-        return _shap_explain_joint_linear_model_cate(self.model_final, X, self._d_t, self._d_y,
-                                                     self.fit_cate_intercept,
+        return _shap_explain_joint_linear_model_cate(self.model_final_, X, self._d_t, self._d_y,
+                                                     self.bias_part_of_coef,
                                                      feature_names=feature_names, treatment_names=treatment_names,
                                                      output_names=output_names,
                                                      input_names=self._input_names,
@@ -736,9 +741,11 @@ class StatsModelsCateEstimatorMixin(LinearModelFinalCateEstimatorMixin):
     Mixin class that offers `inference='statsmodels'` options to the CATE estimator
     that inherits it.
 
-    Such an estimator must implement a :attr:`model_final` attribute that points
+    Such an estimator must implement a :attr:`model_final_` attribute that points
     to the fitted final :class:`.StatsModelsLinearRegression` object that
-    represents the fitted CATE model.
+    represents the fitted CATE model. Also must implement :attr:`featurizer_` that points
+    to the fitted featurizer and :attr:`bias_part_of_coef` that designates
+    if the intercept is the first element of the :attr:`model_final_` coefficient.
     """
 
     def _get_inference_options(self):
@@ -771,7 +778,7 @@ def _get_inference_options(self):
 
     @property
     def feature_importances_(self):
-        return self.model_final.feature_importances_
+        return self.model_final_.feature_importances_
 
 
 class LinearModelFinalCateEstimatorDiscreteMixin(BaseCateEstimator):
@@ -822,7 +829,7 @@ def intercept_(self, T):
         -------
         intercept: float or (n_y,) array like
         """
-        if not self.fit_cate_intercept:
+        if not self.fit_cate_intercept_:
             raise AttributeError("No intercept was fitted!")
         _, T = self._expand_treatments(None, T)
         ind = inverse_onehot(T).item() - 1
@@ -980,7 +987,7 @@ class StatsModelsCateEstimatorDiscreteMixin(LinearModelFinalCateEstimatorDiscret
     Mixin class that offers `inference='statsmodels'` options to the CATE estimator
     that inherits it.
 
-    Such an estimator must implement a :attr:`model_final` attribute that points
+    Such an estimator must implement a :attr:`model_final_` attribute that points
     to a :class:`.StatsModelsLinearRegression` object that is cloned to fit
     each discrete treatment target CATE model and a :attr:`fitted_models_final` attribute
     that returns the list of fitted final models that represent the CATE for each categorical treatment.