allow missing in X for some ortholearner subclasses

Signed-off-by: Fabio Vera <fabiovera@microsoft.com>

econml/_ortho_learner.py

@@ -298,6 +298,9 @@ class _OrthoLearner(TreatmentExpansionMixin, LinearCateEstimator):
         How to aggregate the nuisance value for each sample across the `mc_iters` monte carlo iterations of
         cross-fitting.
 
+    enable_missing: list, default ['W']
+        Which data arrays to allow missing values for. Currently only supports missing values for W, X.
+
     Examples
     --------
 
@@ -434,7 +437,7 @@ def _gen_ortho_learner_model_final(self):
     def __init__(self, *,
                  discrete_treatment, treatment_featurizer,
                  discrete_instrument, categories, cv, random_state,
-                 mc_iters=None, mc_agg='mean'):
+                 mc_iters=None, mc_agg='mean', enable_missing=None):
         self.cv = cv
         self.discrete_treatment = discrete_treatment
         self.treatment_featurizer = treatment_featurizer
@@ -443,6 +446,7 @@ def __init__(self, *,
         self.categories = categories
         self.mc_iters = mc_iters
         self.mc_agg = mc_agg
+        self._enable_missing = enable_missing or []
         super().__init__()
 
     @abstractmethod
@@ -605,9 +609,10 @@ def fit(self, Y, T, *, X=None, W=None, Z=None, sample_weight=None, freq_weight=N
         assert not (self.discrete_treatment and self.treatment_featurizer), "Treatment featurization " \
             "is not supported when treatment is discrete"
         if check_input:
-            Y, T, X, Z, sample_weight, freq_weight, sample_var, groups = check_input_arrays(
-                Y, T, X, Z, sample_weight, freq_weight, sample_var, groups)
-            W, = check_input_arrays(W, force_all_finite='allow-nan')
+            Y, T, Z, sample_weight, freq_weight, sample_var, groups = check_input_arrays(
+                Y, T, Z, sample_weight, freq_weight, sample_var, groups)
+            X, = check_input_arrays(X, force_all_finite='allow-nan' if 'X' in self._enable_missing else True)
+            W, = check_input_arrays(W, force_all_finite='allow-nan' if 'W' in self._enable_missing else True)
             self._check_input_dims(Y, T, X, W, Z, sample_weight, freq_weight, sample_var, groups)
 
         if not only_final:

econml/dml/_rlearner.py

@@ -272,15 +272,16 @@ def _gen_rlearner_model_final(self):
     """
 
     def __init__(self, *, discrete_treatment, treatment_featurizer, categories,
-                 cv, random_state, mc_iters=None, mc_agg='mean'):
+                 cv, random_state, mc_iters=None, mc_agg='mean', enable_missing=None):
         super().__init__(discrete_treatment=discrete_treatment,
                          treatment_featurizer=treatment_featurizer,
                          discrete_instrument=False,  # no instrument, so doesn't matter
                          categories=categories,
                          cv=cv,
                          random_state=random_state,
                          mc_iters=mc_iters,
-                         mc_agg=mc_agg)
+                         mc_agg=mc_agg,
+                         enable_missing=enable_missing)
 
     @abstractmethod
     def _gen_model_y(self):

econml/dml/causal_forest.py

@@ -601,7 +601,8 @@ def __init__(self, *,
                  subforest_size=4,
                  n_jobs=-1,
                  random_state=None,
-                 verbose=0):
+                 verbose=0,
+                 enable_missing=False):
 
         # TODO: consider whether we need more care around stateful featurizers,
         #       since we clone it and fit separate copies
@@ -636,7 +637,8 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=['W'] if enable_missing else None)
 
     def _get_inference_options(self):
         options = super()._get_inference_options()
@@ -738,7 +740,7 @@ def tune(self, Y, T, *, X=None, W=None,
         """
         from ..score import RScorer  # import here to avoid circular import issue
         Y, T, X, sample_weight, groups = check_input_arrays(Y, T, X, sample_weight, groups)
-        W, = check_input_arrays(W, force_all_finite='allow-nan')
+        W, = check_input_arrays(W, force_all_finite='allow-nan' if self._enable_missing else True)
 
         if params == 'auto':
             params = {

econml/dml/dml.py

@@ -108,8 +108,9 @@ def __init__(self, model_final, fit_cate_intercept, featurizer, use_weight_trick
         else:
             self._fit_cate_intercept = fit_cate_intercept
             if self._fit_cate_intercept:
+                # data is already validated at initial fit time
                 add_intercept_trans = FunctionTransformer(add_intercept,
-                                                          validate=True)
+                                                          validate=False)
                 if featurizer:
                     self._featurizer = Pipeline([('featurize', self._original_featurizer),
                                                  ('add_intercept', add_intercept_trans)])
@@ -466,7 +467,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         # TODO: consider whether we need more care around stateful featurizers,
         #       since we clone it and fit separate copies
         self.fit_cate_intercept = fit_cate_intercept
@@ -481,7 +483,8 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=['X', 'W'] if enable_missing else None)
 
     def _gen_featurizer(self):
         return clone(self.featurizer, safe=False)
@@ -692,7 +695,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         super().__init__(model_y=model_y,
                          model_t=model_t,
                          model_final=None,
@@ -705,7 +709,9 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state,)
+                         random_state=random_state,
+                         enable_missing=enable_missing)
+        self._enable_missing = ['W'] if enable_missing else []  # override super's default, which is ['X', 'W']
 
     def _gen_model_final(self):
         return StatsModelsLinearRegression(fit_intercept=False)
@@ -932,7 +938,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.alpha = alpha
         self.n_alphas = n_alphas
         self.alpha_cov = alpha_cov
@@ -952,7 +959,9 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=enable_missing)
+        self._enable_missing = ['W'] if enable_missing else []  # override super's default, which is ['X', 'W']
 
     def _gen_model_final(self):
         return MultiOutputDebiasedLasso(alpha=self.alpha,
@@ -1139,7 +1148,8 @@ def __init__(self, model_y='auto', model_t='auto',
                  bw=1.0,
                  cv=2,
                  mc_iters=None, mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.dim = dim
         self.bw = bw
         super().__init__(model_y=model_y,
@@ -1153,7 +1163,9 @@ def __init__(self, model_y='auto', model_t='auto',
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=enable_missing)
+        self._enable_missing = ['W'] if enable_missing else []  # override super's default, which is ['X', 'W']
 
     def _gen_model_final(self):
         return ElasticNetCV(fit_intercept=False, random_state=self.random_state)
@@ -1326,7 +1338,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
 
         # TODO: consider whether we need more care around stateful featurizers,
         #       since we clone it and fit separate copies
@@ -1340,7 +1353,8 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=['X', 'W'] if enable_missing else None)
 
     def _get_inference_options(self):
         # add blb to parent's options

econml/dr/_drlearner.py

@@ -409,7 +409,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.model_propensity = clone(model_propensity, safe=False)
         self.model_regression = clone(model_regression, safe=False)
         self.model_final = clone(model_final, safe=False)
@@ -423,7 +424,8 @@ def __init__(self, *,
                          treatment_featurizer=None,  # treatment featurization not supported with discrete treatment
                          discrete_instrument=False,  # no instrument, so doesn't matter
                          categories=categories,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=['X', 'W'] if enable_missing else None)
 
     # override only so that we can exclude treatment featurization verbiage in docstring
     def const_marginal_effect(self, X=None):
@@ -871,7 +873,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.fit_cate_intercept = fit_cate_intercept
         super().__init__(model_propensity=model_propensity,
                          model_regression=model_regression,
@@ -883,7 +886,9 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=enable_missing)
+        self._enable_missing = ['W'] if enable_missing else []  # override super's default, which is ['X', 'W']
 
     def _gen_model_final(self):
         return StatsModelsLinearRegression(fit_intercept=self.fit_cate_intercept)
@@ -1149,7 +1154,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.fit_cate_intercept = fit_cate_intercept
         self.alpha = alpha
         self.n_alphas = n_alphas
@@ -1168,7 +1174,9 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=enable_missing)
+        self._enable_missing = ['W'] if enable_missing else []  # override super's default, which is ['X', 'W']
 
     def _gen_model_final(self):
         return DebiasedLasso(alpha=self.alpha,
@@ -1432,7 +1440,8 @@ def __init__(self, *,
                  subforest_size=4,
                  n_jobs=-1,
                  verbose=0,
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.n_estimators = n_estimators
         self.max_depth = max_depth
         self.min_samples_split = min_samples_split
@@ -1456,7 +1465,9 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=enable_missing)
+        self._enable_missing = ['W'] if enable_missing else []  # override super's default, which is ['X', 'W']
 
     def _gen_model_final(self):
         return RegressionForest(n_estimators=self.n_estimators,

econml/iv/dml/_dml.py

@@ -350,7 +350,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.model_y_xw = clone(model_y_xw, safe=False)
         self.model_t_xw = clone(model_t_xw, safe=False)
         self.model_t_xwz = clone(model_t_xwz, safe=False)
@@ -366,7 +367,8 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=['W'] if enable_missing else None)
 
     def _gen_featurizer(self):
         return clone(self.featurizer, safe=False)
@@ -1149,7 +1151,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.model_y_xw = clone(model_y_xw, safe=False)
         self.model_t_xw = clone(model_t_xw, safe=False)
         self.model_t_xwz = clone(model_t_xwz, safe=False)
@@ -1163,7 +1166,8 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=['X', 'W'] if enable_missing else None)
 
     def _gen_featurizer(self):
         return clone(self.featurizer, safe=False)
@@ -1540,7 +1544,8 @@ def __init__(self, *,
                  cv=2,
                  mc_iters=None,
                  mc_agg='mean',
-                 random_state=None):
+                 random_state=None,
+                 enable_missing=False):
         self.model_y_xw = clone(model_y_xw, safe=False)
         self.model_t_xw = clone(model_t_xw, safe=False)
         self.model_t_xwz = clone(model_t_xwz, safe=False)
@@ -1553,7 +1558,8 @@ def __init__(self, *,
                          cv=cv,
                          mc_iters=mc_iters,
                          mc_agg=mc_agg,
-                         random_state=random_state)
+                         random_state=random_state,
+                         enable_missing=['X', 'W'] if enable_missing else None)
 
     def _gen_featurizer(self):
         return clone(self.featurizer, safe=False)