Confused about dimensionality of nuisance model predictions

[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).

When I fit the model with the default n_rep=1 and extract the predictions using dml_model.predictions the dimensionality for dml_model.predictions['ml_m'] is (200, 1, 50) which makes sense (i.e., 200 samples, 1 repetition, 50 treatment variables), but the dimensionality for dml_model.predictions['ml_l'] is the same, and I'm confused by this because there's only one outcome variable.

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)

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

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()

# dml_model.predictions['ml_m'].shape
# (200, 1, 50)

# dml_model.predictions['ml_l'].shape
# (200, 1, 50)

[OliverSchacht]

Hello @RandallJEllis

thanks for your interesting question.
By design and for consistency, DoubleML iterates over all treatments and fits all nuisance models for each treatment. This happens even if a nuisance target is independent of the individual treatments. This explains the dimensionality.

Since you use data = DoubleMLData.from_arrays(x=Z, y=Y, d=T) the default value use_other_treat_as_covariate=True is set and thus, for each treatment, all other 49 treatments join the set of covariates when estimating ml_l and ml_m.

If you'd like to set use_other_treat_as_covariate=False and avoid recalculating ml_l, you can use the doubleml samples to create one set of predictions, broadcast them on all 50 treatments and hand them over as external predictions.

from sklearn.model_selection import cross_val_predict

l_hat_predict = cross_val_predict(estimator=ml_y_z,
                                  X=data.x, 
                                  y=data.y,
                                  cv=dml_model.smpls[0])

l_hat_dict = {f"d{i+1}": {"ml_l": l_hat_predict.reshape(-1,1)} for i in range(50)}

dml_model.fit(external_predictions=l_hat_dict)

Hope this helps!

[RandallJEllis]

Thanks so much!