Predicting the treatment in DML

I'm trying to use the treatment models in DML to get a prediction of the treatment. I can access the models easily with est.models_t

However, it's not trivial to go from those models to a predictions of the treatment. There are some internal transformations on the features, so just passing the original X to the treatment models won't work.

# fits DML with 10 features, one continous treatment
features = ["A", "B", ....]
treatment = "T" # continuous
target = "Y" # continuous

est = LinearDML(model_y =LGBMRegressor(**PARAMS), model_t=LGBMRegressor(**PARAMS), linear_first_stages= False)
est.fit(Y=train_data[target], T=train_data[treatment], X=train_data[features], W=train_data[features])

t_model = est.models_t[0].predict(train_data[features])

ValueError: Number of features of the model must match the input. Model n_features_ is 20 and input n_features is 10 

It would be awesome to have a method that makes predictions using the treatment and outcomes models.

def predict_t(self, X):
    return np.mean([m_t.predict(X) for m_t in self.models_t] axis=1)

Start at LinearDML.fit and the est.models_t entry point, then trace the internal feature transformations applied before the treatment models are used. Done means a public prediction method can accept the original X and return treatment predictions consistently with the fitted models.