CausalForestDML does not accept my inputs with NaN during training, even with a featurizer that deals with NaN

[fburgaud]

Here is a code to see the problem:

import numpy as np
import pandas as pd
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from econml.dml import CausalForestDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.model_selection import train_test_split

# Step 1: Create synthetic dataset
np.random.seed(0)
n = 300
X = pd.DataFrame({
    'age': np.random.normal(40, 10, size=n),
    'income': np.random.normal(50000, 15000, size=n),
    'education_years': np.random.normal(16, 2, size=n)
})

# Introduce missing values in 'income'
missing_idx = np.random.choice(n, size=30, replace=False)
X.loc[missing_idx, 'income'] = np.nan

# Binary treatment and outcome
T = np.random.binomial(1, 0.5, size=n)
Y = 1000 + 500*T + 20*X['age'] + np.random.normal(0, 1000, size=n)

# Step 2: Featurizer to handle missing data
featurizer = Pipeline([
    ('imputer', SimpleImputer(strategy='mean')),
    ('scaler', StandardScaler())
])

# Step 3: CausalForestDML model
est = CausalForestDML(
    model_t=RandomForestClassifier(),
    model_y=RandomForestRegressor(),
    featurizer=featurizer,
    discrete_treatment=True,
    random_state=0
)

# Step 4: Fit the model
est.fit(Y, T, X=X)

Unfortunately this throws an error:

ValueError: Input contains NaN.

The issue appears to be that during fit(), CausalForestDML runs the following:

X, = check_input_arrays(X, force_all_finite='allow-nan' if 'X' in self._gen_allowed_missing_vars() else True)

Which isn't really right because it should rather check the featurized X which is what is ultimately used for fitting anyways, not X itself.

[fburgaud]

@fverac Tagging you as this is in line with some other discussions on input verifications that are more restrictive than what one might like.

[fverac]

I see what you are getting at.

Note, the featurizer is only applied in the final stage, so even if input checking were completely disabled on the econml side, the repro code you shared would still fail because the nuisance models you supply (random forests) don't handle missing values natively, and they would receive the unfeaturized X.

A slightly more pointed repro code snippet would be:

import numpy as np
import pandas as pd
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from econml.dml import CausalForestDML, NonParamDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.model_selection import train_test_split
from lightgbm import LGBMRegressor, LGBMClassifier

# Step 1: Create synthetic dataset
np.random.seed(0)
n = 300
X = pd.DataFrame({
    'age': np.random.normal(40, 10, size=n),
    'income': np.random.normal(50000, 15000, size=n),
    'education_years': np.random.normal(16, 2, size=n)
})

# Introduce missing values in 'income'
missing_idx = np.random.choice(n, size=30, replace=False)
X.loc[missing_idx, 'income'] = np.nan

# Binary treatment and outcome
T = np.random.binomial(1, 0.5, size=n)
Y = 1000 + 500*T + 20*X['age'] + np.random.normal(0, 1000, size=n)

# Step 2: Featurizer to handle missing data
featurizer = Pipeline([
    ('imputer', SimpleImputer(strategy='mean')),
    ('scaler', StandardScaler())
])

# Step 3: CausalForestDML model
est = CausalForestDML(
    model_t=LGBMClassifier(), #lgbm models can handle missing values natively
    model_y=LGBMRegressor(),
    featurizer=featurizer,
    discrete_treatment=True,
    random_state=0,
    allow_missing=True # note allow missing only disables input checking for W in CausalForestDML
)

# Step 4: Fit the model
est.fit(Y, T, X=X)

Here we use LGBM models for nuisances since they handle missing values natively. So ideally this code should work but it doesn't.

A few workarounds come to mind.

1. You could featurize your model before passing to econml. But then your nuisance models would also be fitting on a featurized X. Not sure if that is preferred.
2. You could do a little hack to get the models to allow missing, like below:

import numpy as np
import pandas as pd
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from econml.dml import CausalForestDML, NonParamDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.model_selection import train_test_split
from lightgbm import LGBMRegressor, LGBMClassifier

# Step 1: Create synthetic dataset
np.random.seed(0)
n = 300
X = pd.DataFrame({
    'age': np.random.normal(40, 10, size=n),
    'income': np.random.normal(50000, 15000, size=n),
    'education_years': np.random.normal(16, 2, size=n)
})

# Introduce missing values in 'income'
missing_idx = np.random.choice(n, size=30, replace=False)
X.loc[missing_idx, 'income'] = np.nan

# Binary treatment and outcome
T = np.random.binomial(1, 0.5, size=n)
Y = 1000 + 500*T + 20*X['age'] + np.random.normal(0, 1000, size=n)

# Step 2: Featurizer to handle missing data
featurizer = Pipeline([
    ('imputer', SimpleImputer(strategy='mean')),
    ('scaler', StandardScaler())
])

# Step 3: CausalForestDML model
est = CausalForestDML(
    model_t=LGBMClassifier(),
    model_y=LGBMRegressor(),
    featurizer=featurizer,
    discrete_treatment=True,
    random_state=0,
    allow_missing=True
)

def new_gen_allowed_missing_vars(self):
    # Override to allow missing values in X
    return ['X', 'W']
CausalForestDML._gen_allowed_missing_vars = new_gen_allowed_missing_vars

# Step 4: Fit the model
est.fit(Y, T, X=X)

The second hack works most like your original intention, but it's still a hack so not very clean looking.

Longer term solution would be to potentially expand the accepted values that the allow_missing param takes. In addition to True/False, maybe they could take an additional kwarg like 'disable' which allows missing in 'X', 'W' always. Or they could accept a list of characters like ['X', 'W']. Have to think about this one, but hope the above temporary solutions work for now.

[fburgaud]

Thanks for the hack, will keep it in mind for the time being.

Just a few notes. Random Forest does handle missing values since sklearn v1.4.

Other thing I'd point out is that DML (with Lasso for example) would not have failed in my code, while CausalForestDML and LinearDML do, which is inconsistent in behavior.

[fverac]

Good to know about random forests.

As for the inconsistency, the way the allow_missing feature is designed currently is to only allow missing values for X when a final model can be passed that handles missing values.

CausalForestDML and LinearDML have a model_final that cannot handle missing values, and since they will need to fit on X (or some featurization of X), that is why they only accept missing in W when allow_missing is True.

On the other hand, a user can pass a model_final to DML that does handle missing values, which is why missing values in X are allowed for that class.