Error when making predictions with nuisance models

[RandallJEllis]

I am running simulations with 200 samples, 2 covariates to orthogonalize with respect to (Z), 50 treatment variables I want to orthogonalize (T), and 1 binary outcome I want to orthogonalize. I am running DoubleMLPLR with two nuisance models (i.e., ml_m argument for T|Z and ml_l argument for Y|Z).

What I'm looking to do is:

1. Split the data into train (80%) and test (20%)
2. Run DML on the training set to calculate coefficients
3. Use the fitted nuisance models to orthogonalize T and Y net Z in the test set
4. Use the calculated coefficients to make predictions on the orthogonalized T to predict orthogonalized Y in the test set

However, when I use store_models=True and call the fitted nuisance models, I get an error about the number of features in the data. Specifically, it says the training data has 51 features in it, when I have 2 covariates and 50 treatments. Code and output below. Any help is appreciated!

import numpy as np
import pandas as pd
from sklearn.model_selection import KFold
from lightgbm import LGBMClassifier, LGBMRegressor
from joblib import Parallel, delayed
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.multioutput import MultiOutputRegressor
from doubleml import DoubleMLData, DoubleMLPLR
from lightgbm import LGBMRegressor, LGBMClassifier
from sklearn.model_selection import KFold

n_samples = 200
n_treatments = 50

# Generate mock data
np.random.seed(42)

Z = pd.DataFrame({
    'covar1': np.random.uniform(50, 90, n_samples),
    'covar2': np.random.choice([0, 1], size=n_samples),  # Categorical variable
                     })

n_covars = Z.shape[1]

# Create T to be correlated with Z
# Generate random weights for correlation
weights = np.random.randn(n_covars, n_treatments)

# Create T to be correlated with Z
T = pd.DataFrame((Z.values @ weights + np.random.randn(n_samples, n_treatments)) / 10, columns=[f"P{i}" for i in range(n_treatments)])

# Create Y to be correlated with Z and T
# Generate random weights for correlation
weights_Z = np.random.randn(n_covars)
weights_T = np.random.randn(n_treatments)

# Compute continuous Y
Y_continuous = Z.values @ weights_Z + T.values @ weights_T + np.random.randn(n_samples)

# Binarize Y at a cutpoint to make ~10% positives
cutpoint = np.percentile(Y_continuous, 90)  # Top 10% as positives
Y = (Y_continuous > cutpoint).astype(int).reshape(-1, 1)

train_idx, test_idx = train_test_split(np.arange(len(Y)), test_size=0.2, stratify=Y, random_state=42)

# Split the data into training and test sets
Z_train, Z_test = Z.iloc[train_idx], Z.iloc[test_idx]
T_train, T_test = T.iloc[train_idx], T.iloc[test_idx]
Y_train, Y_test = Y[train_idx], Y[test_idx]

# Setup: outcome Y, covariates Z, treatment T 
data = DoubleMLData.from_arrays(x=Z_train, y=Y_train, d=T_train)

ml_t_z = LGBMRegressor()
ml_y_z = LGBMClassifier()

dml_model = DoubleMLPLR(data, 
                        ml_l=ml_y_z, 
                        ml_m=ml_t_z, 
                        n_folds=5
                        )
dml_model.fit(store_models=True)

dml_model.models['ml_l']['d1'][0][0].predict(Z_test)

LightGBMError: The number of features in data (2) is not the same as it was in training data (51).
You can set ``predict_disable_shape_check=true`` to discard this error, but please be aware what you are doing.

dml_model.models['ml_m']['d1'][0][0].predict(Z_test)

LightGBMError: The number of features in data (2) is not the same as it was in training data (51).
You can set ``predict_disable_shape_check=true`` to discard this error, but please be aware what you are doing.

[OliverSchacht]

Hi @RandallJEllis - thanks for pointing this out.

I think this question is also related to my answer in #342 which I have to adjust a bit. Sorry for the confusion.

When you do:

data = DoubleMLData.from_arrays(x=Z_train, y=Y_train, d=T_train)

You use the default value for use_other_treat_as_covariate, which is True (I thought it was False 😅) .

So in fact, for each treatment, when estimating E[ Y | Z* ] and E[ T | Z* ], Z* contains the two covariates Z as well as the 49 other treatment columns. This is why LGBM puts out the dimensionality error in the code provided. If you set it to

data = DoubleMLData.from_arrays(x=Z_train, y=Y_train, d=T_train, use_other_treat_as_covariate=False)

the error disappears.

Hope this clears the confusion!

[RandallJEllis]

Thanks again, Oliver!