Issue with #791: Missing data allowed when training, but not inferencing

[esbraun]

#791 ought to allow missing data when inferencing if allow_missing = True, but currently an error is thrown when the inference data includes missing data. RE: below for a working example based on the "Metalearner Examples" notebook for a demonstration of the issue.

https://colab.research.google.com/drive/1QT-IdwRWRUiQNDwqJ0EXPXSftDuIg3uE?usp=sharing

[Ftrejo23]

I'm seeing the same issue

[fverac]

Thanks for raising this issue. In the original PR we originally only considered missing data during training, but it seems there is interest in extending this to inferencing as well.

[erasedcitizen11]

Hi there, just wondering if there any update on this as running into the same problems (doesn't allow missing values when computing CATE/ATE).

[esbraun]

@erasedcitizen11 I created a fork that disables the missing value checks which is a requirement for our use cases. Use at your own risk: https://github.com/esbraun/EconML/tree/main

[pgeertsema]

@EconML team. Great product!

Any progress in remediating this issue? It appears that scikit-learn's default input validation checks remain in place even when allow_missing=True is set in DML() and the base learners support missing values (such as LightGBM). This is inconsistent with the documentation (at https://econml.azurewebsites.net/_autosummary/econml.dml.DML.html, see allow_missing (bool)). A self-contained minimal working example follows:

import numpy as np
import scipy.special as sc
from econml.dml import DML
from econml.sklearn_extensions.linear_model import StatsModelsLinearRegression
from lightgbm import LGBMRegressor, LGBMClassifier

np.random.seed(123)
X = np.random.normal(size=(1000, 5))
T = np.random.binomial(1, sc.expit(X[:, 0]))

mask = np.random.rand(*X.shape) < 0.05
X[mask] = np.nan

y = (1 + .5*X[:, 0]) * T + X[:, 0] + np.random.normal(size=(1000,))
est = DML(
    # both base learner models support missing values out of the box
    model_y=LGBMRegressor(),
    model_t=LGBMClassifier(),
    model_final=StatsModelsLinearRegression(fit_intercept=False),
    # should allow discrete treatment
    discrete_treatment=True,
    allow_missing=True,
)
est.fit(y, T, X=X, W=None)

It terminates with

Traceback (most recent call last):
  File "c:\Users\Paul\Dropbox\DML\MLcode\another test.py", line 24, in <module>
    est.fit(y, T, X=X, W=None)
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\dml\dml.py", line 589, in fit
    return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\dml\_rlearner.py", line 415, in fit
    return super().fit(Y, T, X=X, W=W,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\_cate_estimator.py", line 131, in call
    m(self, Y, T, *args, **kwargs)
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\_ortho_learner.py", line 745, in fit
    Y, T, Z, sample_weight, freq_weight, sample_var, groups = check_input_arrays(
                                                              ^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\utilities.py", line 578, in check_input_arrays
    new_arg = check_array(arg, dtype=dtype, ensure_2d=False, accept_sparse=True,
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\sklearn\utils\validation.py", line 1064, in check_array
    _assert_all_finite(
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\sklearn\utils\validation.py", line 123, in _assert_all_finite
    _assert_all_finite_element_wise(
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\sklearn\utils\validation.py", line 172, in _assert_all_finite_element_wise
    raise ValueError(msg_err)
ValueError: Input contains NaN.

@esbraun: Thanks for the workaround; this is helpful. In the longer term it would be better if we could remediate the issue in the main branch.

[kbattocchi]

@pgeertsema we allow missing values in W and X, but your DGP also produces missing values in Y, which we never allow. You can just filter those rows out if you want to fit on the remaining data.

[esbraun]

@kbattocchi To clarify - missing data is still not yet allowed when inferencing? Thank you!

[pgeertsema]

@kbattocchi Yes, of course if there are missing values in Y, then I can see why it would not work. As noted by @esbraun, the issue is that the second stage (inference) fails even if there are no missing values in Y.

The code snippet below has no missing values in Y and base models that can handle missing values in X, but still fails:

import numpy as np
import scipy.special as sc
from econml.dml import DML
from econml.sklearn_extensions.linear_model import StatsModelsLinearRegression
from lightgbm import LGBMRegressor, LGBMClassifier

np.random.seed(123)
X = np.random.normal(size=(1000, 5))
T = np.random.binomial(1, sc.expit(X[:, 0]))

# mo missing values in Y
y = (1 + .5*X[:, 0]) * T + X[:, 0] + np.random.normal(size=(1000,))

# missing values in X only
mask = np.random.rand(*X.shape) < 0.05
X[mask] = np.nan

est = DML(
    # both base learner models support missing values out of the box
    model_y=LGBMRegressor(),
    model_t=LGBMClassifier(),
    model_final=StatsModelsLinearRegression(fit_intercept=False),
    # should allow discrete treatment
    discrete_treatment=True,
    allow_missing=True,
)
est.fit(y, T, X=X, W=None)

Running the above yields:

Traceback (most recent call last):
  File "c:\Users\Paul\Dropbox\DML\MLcode\EconML error.py", line 27, in <module>
    est.fit(y, T, X=X, W=None)
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\dml\dml.py", line 589, in fit
    return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\dml\_rlearner.py", line 415, in fit
    return super().fit(Y, T, X=X, W=W,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\_cate_estimator.py", line 131, in call
    m(self, Y, T, *args, **kwargs)
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\_ortho_learner.py", line 877, in fit
    self._fit_final(Y=Y,
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\_ortho_learner.py", line 980, in _fit_final
    self._ortho_learner_model_final.fit(Y, T, **filter_none_kwargs(X=X, W=W, Z=Z,
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\dml\_rlearner.py", line 94, in fit
    self._model_final.fit(X, T, T_res, Y_res, sample_weight=sample_weight,
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\dml\dml.py", line 163, in fit
    self._model.fit(fts, Y_res, **filtered_kwargs)
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\sklearn_extensions\linear_model.py", line 1794, in fit  
    X, y, freq_weight, sample_var = self._check_input(X, y, sample_weight, freq_weight, sample_var)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\sklearn_extensions\linear_model.py", line 1706, in _check_input
    X, y, sample_weight, freq_weight, sample_var = check_input_arrays(
                                                   ^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\econml\utilities.py", line 578, in check_input_arrays
    new_arg = check_array(arg, dtype=dtype, ensure_2d=False, accept_sparse=True,
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\sklearn\utils\validation.py", line 1064, in check_array        
    _assert_all_finite(
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\sklearn\utils\validation.py", line 123, in _assert_all_finite  
    _assert_all_finite_element_wise(
  File "C:\Users\Paul\anaconda3\envs\dml\Lib\site-packages\sklearn\utils\validation.py", line 172, in _assert_all_finite_element_wise
    raise ValueError(msg_err)
ValueError: Input contains NaN.

Sorry for the earlier confusion; hopefully this makes clear where we see the inconsistency with the documentation.

[jhn-nt]

I'd love to see a fix for this as well, specially for X. I usually defer imputation to a transformer object before an estimator in a scikit-learn pipeline, this makes things cleaner at validation time.
Is there any plan to extend support for missing values on X at inference time?

Best
Giovanni

[fburgaud]

Any updates in this? It's very weird to allow a model to train on null values, but not let it predict with inputs with null values.
Sometimes, I even have a preprocessor to deal with the null values (e.g. for a LinearDML), but it still thinks it's bad.
That's not great for usability.

[fverac]

Late to this conversation, but the reason @pgeertsema's repro code fails is because the model_final passed to DML does not natively handle missing values.

If you use a different model final that can handle missing values, the code snippet runs successfully. See below. In this case I just made model_final an sklearn pipeline with a SimpleImputer as the first step.

import numpy as np
import scipy.special as sc
from econml.dml import DML
from econml.sklearn_extensions.linear_model import StatsModelsLinearRegression
from lightgbm import LGBMRegressor, LGBMClassifier
from sklearn.pipeline import make_pipeline
from sklearn.impute import SimpleImputer


np.random.seed(123)
X = np.random.normal(size=(1000, 5))
T = np.random.binomial(1, sc.expit(X[:, 0]))

# mo missing values in Y
y = (1 + .5*X[:, 0]) * T + X[:, 0] + np.random.normal(size=(1000,))

# missing values in X only
mask = np.random.rand(*X.shape) < 0.05
X[mask] = np.nan

est = DML(
    # both base learner models support missing values out of the box
    model_y=LGBMRegressor(),
    model_t=LGBMClassifier(),
    model_final=make_pipeline(SimpleImputer(strategy='mean'), StatsModelsLinearRegression()),
    # should allow discrete treatment
    discrete_treatment=True,
    allow_missing=True,
)
est.fit(y, T, X=X, W=None)

Though the problem with the code I shared above is that est.summary() bugs out and complains that 'Pipeline' object has no attribute 'coef', which is true. Instead, a better way to pass a Pipeline model_final would be by wrapping it in a class as follows:

import numpy as np
import scipy.special as sc
from econml.dml import DML
from econml.sklearn_extensions.linear_model import StatsModelsLinearRegression
from lightgbm import LGBMRegressor, LGBMClassifier
from sklearn.pipeline import make_pipeline
from sklearn.impute import SimpleImputer


np.random.seed(123)
X = np.random.normal(size=(1000, 5))
T = np.random.binomial(1, sc.expit(X[:, 0]))

# mo missing values in Y
y = (1 + .5*X[:, 0]) * T + X[:, 0] + np.random.normal(size=(1000,))

# missing values in X only
mask = np.random.rand(*X.shape) < 0.05
X[mask] = np.nan

# class the correctly exposes attributes from the final stage of an sklearn pipeline
# for use in post-econml-fit methods like summary()
class ParametricModelFinalForMissing:
    def __init__(self, model_final):
        self.model_final = model_final

    def fit(self, *args, **kwargs):
        self.model_final.fit(*args, **kwargs)

        self.coef_ = self.model_final[-1].coef_
        self.intercept_ = self.model_final[-1].intercept_
        self.intercept_stderr_ = self.model_final[-1].intercept_stderr_
        self.coef_stderr_ = self.model_final[-1].coef_stderr_
        return self

    def predict(self, *args, **kwargs):
        return self.model_final.predict(*args, **kwargs)

    # assumes an sklearn pipeline by indexing -1 of model_final
    def prediction_stderr(self, *args, **kwargs):
        return self.model_final[-1].prediction_stderr(*args, **kwargs)

est = DML(
    # both base learner models support missing values out of the box
    model_y=LGBMRegressor(),
    model_t=LGBMClassifier(),
    model_final=ParametricModelFinalForMissing(make_pipeline(SimpleImputer(strategy='mean'), StatsModelsLinearRegression())),
    # should allow discrete treatment
    discrete_treatment=True,
    allow_missing=True,
)
est.fit(y, T, X=X, W=None)
print(est.summary())

[fburgaud]

@fverac Have you tried this for inference though? I've used final models that accept nulls, but the inference step failed because there is a table check step at the beginning that errors out the process before it even tries to use the models.