Ensuring Fully Time-Aware Cross-Validation with FLAML & DoubleML to Prevent Data Leakage in Time Series

[paul-jdfagan]

Hello,
Huge fan of the package. I'm trying to apply it and FLAML to time series data but Im getting tripped up ensuring I've a fully time-aware CV fit and avoid data leakage etc.

Below is the base code i'm using. I use "split_type="time" for the FlamlRegressorDoubleML, but the final DoubleML Cross-Fitting will not be time-aware.

dml_plr = dml.DoubleMLPLR(
    dml_data,
    ml_l=ml_l,
    ml_m=ml_m,
    **n_folds=5,
    n_rep=5,**
    score="partialling out",
    **draw_sample_splitting=True**
) 

My ask: I would appreciate any advice for feeding time‐stamped data into a standard K-Fold splitter and avoid the data leakage issues. Maybe im deep down a rabbit hole of overthinking it

I believe I'd need to change the automatic splitting 'draw_sample_splitting=False', and do my own TimeSeriesSplit , but previous attempts would not run properly - run into partitions errors. An alternative is I don't respect the final CV time series ordering and just accept that my final model results are not time sensitive. Maybe its good enough that my two nuisance models are. All very confusing...

I'd greatly appreciate any support. Thank you.

Paul

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from flaml import AutoML
import doubleml as dml
import logging
import time

# --- Configure Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 1) Pure temporal transformer
class PureTemporalFeatures(BaseEstimator, TransformerMixin):
    """
    Add only date-based, leak-safe features.
    """
    def __init__(self, date_col="Date"):
        self.date_col = date_col

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        t0 = time.time()
        logger.info("Starting PureTemporalFeatures.transform")
        df = X.copy()
        df[self.date_col] = pd.to_datetime(df[self.date_col])
        df = df.sort_values(self.date_col).reset_index(drop=True)

        # Basic cyclical encodings
        df["month_sin"] = np.sin(2*np.pi * df[self.date_col].dt.month/12)
        df["month_cos"] = np.cos(2*np.pi * df[self.date_col].dt.month/12)
        df["dow_sin"]   = np.sin(2*np.pi * df[self.date_col].dt.weekday/7)
        df["dow_cos"]   = np.cos(2*np.pi * df[self.date_col].dt.weekday/7)
        df["doy_sin"]   = np.sin(2*np.pi * df[self.date_col].dt.dayofyear/365)
        df["doy_cos"]   = np.cos(2*np.pi * df[self.date_col].dt.dayofyear/365)

        # calendar flags
        df["is_weekend"]       = (df[self.date_col].dt.weekday >= 5).astype(int)
        df["is_month_start"]   = (df[self.date_col].dt.day <= 7).astype(int)
        df["is_month_end"]     = (df[self.date_col].dt.day >= 24).astype(int)
        df["is_quarter_start"] = df[self.date_col].dt.month.isin([1,4,7,10]).astype(int)
        df["is_quarter_end"]   = df[self.date_col].dt.month.isin([3,6,9,12]).astype(int)

        # shopping-season flags
        df["is_holiday_season"]           = (
            (df[self.date_col].dt.month==12) |
            ((df[self.date_col].dt.month==11)&(df[self.date_col].dt.day>=15))
        ).astype(int)
        df["is_black_friday_week"]        = (
            (df[self.date_col].dt.month==11)&(df[self.date_col].dt.day>=20)
        ).astype(int)
        df["is_nordstrom_anniversary_sale"] = (
            ((df[self.date_col].dt.month==7)&(df[self.date_col].dt.day>=15)) |
            ((df[self.date_col].dt.month==8)&(df[self.date_col].dt.day<=10))
        ).astype(int)
        df["is_sephora_spring_sale"]      = (
            (df[self.date_col].dt.month==4)&
            df[self.date_col].dt.day.between(1,20)
        ).astype(int)
        df["is_walmart_july_event"]       = (
            (df[self.date_col].dt.month==7)&
            df[self.date_col].dt.day.between(10,20)
        ).astype(int)
        df["is_target_july_event"]        = df["is_walmart_july_event"]
        df["is_amazon_prime_days_window"] = (
            (df[self.date_col].dt.month==7)&
            df[self.date_col].dt.day.between(9,18)
        ).astype(int)

        # simple trend
        df["time_trend"]    = np.arange(len(df))
        df["time_trend_sq"] = df["time_trend"] ** 2

        elapsed = time.time() - t0
        logger.info(f"PureTemporalFeatures.transform completed in {elapsed:.2f}s, added {len(df.columns)-1} features")
        return df.drop(columns=[self.date_col])


# 2) FLAML wrapper for DoubleML
class FlamlRegressorDoubleML:
    _estimator_type = "regressor"
    def __init__(self, time=120, estimator_list=None, metric="rmse", **kwargs):
        self.auto_ml = AutoML(**kwargs)
        self.time = time
        self.estimator_list = estimator_list or ["lgbm","xgboost","histgb","xgb_limitdepth","catboost"]
        self.metric = metric
        self.best_loss_ = None
        self.tuned_model = None

    def set_params(self, **params):
        self.auto_ml.set_params(**params)
        return self

    def get_params(self, deep=True):
        p = self.auto_ml.get_params(deep)
        p.update(time=self.time, estimator_list=self.estimator_list, metric=self.metric)
        return p

    def fit(self, X, y):
        t0 = time.time()
        logger.info(f"FLAML.fit: starting on {X.shape[0]} rows, budget={self.time}s")
        self.auto_ml.fit(
            X, y,
            task="regression",
            time_budget=self.time,
            estimator_list=self.estimator_list,
            metric=self.metric,
            split_type="time",
            verbose=False
        )
        self.tuned_model = self.auto_ml.model.estimator
        try:
            self.best_loss_ = self.auto_ml.best_loss
        except:
            self.best_loss_ = getattr(self.auto_ml._state, "best_loss", None)
        elapsed = time.time() - t0
        logger.info(f"FLAML.fit completed in {elapsed:.2f}s, best_loss_={self.best_loss_}")
        return self

    def predict(self, X):
        return self.tuned_model.predict(X)


# 3) Load & preprocess 
t0_total = time.time()
logger.info("Starting data load & preprocessing")
df = dataset
df["Amount spent (USD)"].fillna(method='ffill', inplace=True)

def geometric_adstock(s, alpha=0.9, lags=7):
    w = alpha ** np.arange(lags)
    return s.rolling(lags, min_periods=1) \
            .apply(lambda x: np.dot(x[::-1], w[:len(x)]), raw=True)

df["adstock_spend"] = geometric_adstock(df["Amount spent (USD)"])
df.dropna(subset=["gmv","adstock_spend"], inplace=True)
df.reset_index(drop=True, inplace=True)
logger.info(f"Data loaded & preprocessed in {time.time() - t0_total:.2f}s — {len(df)} obs")

#  4) Feature engineering & DMLData 
t0 = time.time()
logger.info("Applying PureTemporalFeatures")
df_temp = PureTemporalFeatures(date_col="Date").fit_transform(df)

business = ["10_yr_ir","dj_close_open_ratio","App DAUs (iOS+Android)","posts"]
temp_feats = [
    "month_sin","month_cos","dow_sin","dow_cos","doy_sin","doy_cos",
    "is_weekend","is_month_start","is_month_end",
    "is_quarter_start","is_quarter_end","is_holiday_season","is_black_friday_week",
    "is_nordstrom_anniversary_sale","is_sephora_spring_sale",
    "is_walmart_july_event","is_target_july_event","is_amazon_prime_days_window",
    "time_trend","time_trend_sq"
]
all_x = business + temp_feats

df_model = pd.concat([
    df_temp[temp_feats],
    df[business],
    df[["gmv","adstock_spend"]]
], axis=1).loc[:, lambda d: ~d.columns.duplicated()]

dml_data = dml.DoubleMLData(df_model, y_col="gmv", d_cols="adstock_spend", x_cols=all_x)
logger.info(f"Feature engineering & DMLData creation took {time.time() - t0:.2f}s — {dml_data.n_obs} obs × {len(all_x)} covs")

# 5) Nuisance learners 
t0 = time.time()
logger.info("Setting up nuisance learners")
ml_l = make_pipeline(
    SimpleImputer(strategy="median"),
    StandardScaler(),
    FlamlRegressorDoubleML(time=120, metric="rmse")
)
ml_m = make_pipeline(
    SimpleImputer(strategy="median"),
    StandardScaler(),
    FlamlRegressorDoubleML(time=120, metric="rmse")
)
logger.info(f"Nuisance learners ready in {time.time() - t0:.2f}s")

# 6) DoubleMLPLR fit 
t0 = time.time()
logger.info("Fitting DoubleMLPLR (n_folds=5, n_rep=1)")
dml_plr = dml.DoubleMLPLR(
    dml_data,
    ml_l=ml_l,
    ml_m=ml_m,
    n_folds=5,
    n_rep=10,
    score="partialling out"
)
res = dml_plr.fit(n_jobs_cv=-1, store_models=True)
logger.info(f"DoubleMLPLR fit completed in {time.time() - t0:.2f}s")

#  7) Results & diagnostics 
logger.info("Causal effect estimate:")
logger.info(f"\n{res.summary}")

fl_l = dml_plr.models['ml_l'][0].named_steps['flamlregressordoubleml'].best_loss_
fl_m = dml_plr.models['ml_m'][0].named_steps['flamlregressordoubleml'].best_loss_
logger.info(f"FLAML tuning RMSEs → ml_l: {fl_l:.3f}, ml_m: {fl_m:.3f}")

evals = dml_plr.evaluate_learners()
for learner, scores in evals.items():
    rmse = np.mean([s[0] for s in scores])
    logger.info(f"OOS RMSE ({learner}): {rmse:.2f}")

logger.info(f"Total script time: {time.time() - t0_total:.2f}s")

[SvenKlaassen]

Hi @paul-jdfagan,

generally I would agree that a time-aware CV procedure should be used if you would like to apply DoubleML to time series data.
And i can see that external partitioning might not be very simple to use in this setting.

A possible alternative would be the use of external predictions. This way you would be able to fully control the CV-procedure.

I will try to take a more detailed look at your issue in the upcoming week.
You could also provide me with your current version of the time-relevant sample splitting (and I can try to see what might be the reason for the partitions errors).

[paul-jdfagan]

Thanks @SvenKlaassen for looking into the case. Much appreciated.

Below is an updated flow trying to build around the external_predictions argument. Apologies if its overcomplicated.

• I needed to add a fixed-schema guard as features were being dropped and then not found in testing.
• The latest error is seems to be NaNs despite there being none in the dataset. Somehow the model returned NaN predictions.

Thanks again!

2025-06-02 14:39:06,220 - INFO - Loading data …
2025-06-02 14:39:06,239 - INFO - Data ready: 418 obs
2025-06-02 14:39:06,252 - INFO - Feature engineering & DMLData creation took 0.01s — 418 obs × 24 covs
2025-06-02 14:39:06,253 - INFO - Generating external predictions with TimeSeriesSplit
2025-06-02 14:39:06,256 - INFO - Fold 1: train 0–105 | test 106–209
2025-06-02 14:41:06,271 - INFO - Fold 2: train 0–209 | test 210–313
2025-06-02 14:43:06,354 - INFO - Fold 3: train 0–313 | test 314–417
2025-06-02 14:45:06,427 - INFO - CV finished in 360.17s

AssertionError Traceback (most recent call last)
Cell In[158], line 258
255 preds_m_ext[tei, 0] = robust_m.predict(dml_data.x[tei, :])
257 logger.info("CV finished in %.2fs", time.time() - t0_manual_cv)
--> 258 assert not np.isnan(preds_l_ext).any()
259 assert not np.isnan(preds_m_ext).any()
261 external_predictions_dict = {
262 "adstock_spend": {"ml_l": preds_l_ext, "ml_m": preds_m_ext}
263 }

AssertionError:

import numpy as np
import pandas as pd
import logging
import time
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.model_selection import TimeSeriesSplit
from sklearn.metrics import mean_squared_error
from flaml import AutoML
import doubleml as dml

# Global timers & constants 
T0_TOTAL = time.time()
N_FOLDS_TS = 3
N_REP_MANUAL = 1

# Configure Logging 
logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# 1) Pure temporal transformer (leak‑safe)
class PureTemporalFeatures(BaseEstimator, TransformerMixin):
    """Generate leak‑safe calendar features."""

    def __init__(self, date_col: str = "Date"):
        self.date_col = date_col

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        t0 = time.time()
        df = X.copy()
        if self.date_col not in df.columns:
            raise KeyError(f"Date column '{self.date_col}' not found.")

        df[self.date_col] = pd.to_datetime(df[self.date_col])
        df = df.sort_values(self.date_col).reset_index(drop=True)

        # Cyclical encodings
        df["month_sin"] = np.sin(2 * np.pi * df[self.date_col].dt.month / 12)
        df["month_cos"] = np.cos(2 * np.pi * df[self.date_col].dt.month / 12)
        df["dow_sin"]   = np.sin(2 * np.pi * df[self.date_col].dt.weekday / 7)
        df["dow_cos"]   = np.cos(2 * np.pi * df[self.date_col].dt.weekday / 7)
        df["doy_sin"]   = np.sin(2 * np.pi * df[self.date_col].dt.dayofyear / 365.25)
        df["doy_cos"]   = np.cos(2 * np.pi * df[self.date_col].dt.dayofyear / 365.25)

        # Flags
        df["is_weekend"] = (df[self.date_col].dt.weekday >= 5).astype(int)
        df["is_month_start"] = (df[self.date_col].dt.day <= 7).astype(int)
        df["is_month_end"]   = (df[self.date_col].dt.day >= 24).astype(int)
        df["is_quarter_start"] = df[self.date_col].dt.is_quarter_start.astype(int)
        df["is_quarter_end"]   = df[self.date_col].dt.is_quarter_end.astype(int)

        df["is_holiday_season"] = ((df[self.date_col].dt.month == 12) |
                                     ((df[self.date_col].dt.month == 11) &
                                      (df[self.date_col].dt.day >= 15))).astype(int)
        df["is_black_friday_week"] = ((df[self.date_col].dt.month == 11) &
                                        df[self.date_col].dt.day.between(20, 30)).astype(int)
        df["is_nordstrom_anniversary_sale"] = (((df[self.date_col].dt.month == 7) &
                                                  (df[self.date_col].dt.day >= 15)) |
                                                 ((df[self.date_col].dt.month == 8) &
                                                  (df[self.date_col].dt.day <= 10))).astype(int)
        df["is_sephora_spring_sale"] = ((df[self.date_col].dt.month == 4) &
                                          df[self.date_col].dt.day.between(1, 20)).astype(int)
        df["is_walmart_july_event"]  = ((df[self.date_col].dt.month == 7) &
                                          df[self.date_col].dt.day.between(10, 20)).astype(int)
        df["is_target_july_event"]   = df["is_walmart_july_event"]
        df["is_amazon_prime_days_window"] = ((df[self.date_col].dt.month == 7) &
                                               df[self.date_col].dt.day.between(9, 18)).astype(int)

        # Time trend
        df["time_trend"]    = np.arange(len(df))
        df["time_trend_sq"] = df["time_trend"] ** 2

        return df.drop(columns=[self.date_col])

# 2) Fixed‑schema guard 
class FixedSchema(BaseEstimator, TransformerMixin):
    """Freeze feature schema at fit() and reproduce at transform()."""

    def __init__(self, fill_value: float = 0.0, as_ndarray: bool = True):
        self.fill_value = fill_value
        self.as_ndarray = as_ndarray

    def fit(self, X, y=None):
        if isinstance(X, pd.DataFrame):
            self.columns_ = list(X.columns)
            self.n_features_in_ = len(self.columns_)
        else:
            self.columns_ = None
            self.n_features_in_ = X.shape[1]
        return self

    def transform(self, X):
        if isinstance(X, pd.DataFrame):
            Xf = X.copy()
            for col in set(self.columns_) - set(Xf.columns):
                Xf[col] = self.fill_value
            Xf = Xf[self.columns_]
        else:
            X_arr = np.asarray(X)
            if X_arr.shape[1] < self.n_features_in_:
                pad = np.full((X_arr.shape[0], self.n_features_in_ - X_arr.shape[1]), self.fill_value)
                Xf = np.hstack([X_arr, pad])
            else:
                Xf = X_arr[:, :self.n_features_in_]
        return Xf if self.as_ndarray else pd.DataFrame(Xf, columns=self.columns_)

#3) FLAML regressor wrapper 
class FlamlRegressorDoubleML(BaseEstimator):
    _estimator_type = "regressor"

    def __init__(self, time_budget: int = 60, metric: str = "rmse", estimator_list=None):
        self.auto_ml = AutoML()
        self.time_budget = time_budget
        self.metric = metric
        self.estimator_list = estimator_list or [
            "lgbm", "xgboost", "histgb", "catboost", "histgb", "rf"
        ]

    def fit(self, X, y):
        X_df = pd.DataFrame(X)
        self.auto_ml.fit(
            X_train=X_df,
            y_train=y,
            task="regression",
            time_budget=self.time_budget,
            estimator_list=self.estimator_list,
            metric=self.metric,
            split_type="time",
            verbose=0,
        )
        self.model_ = self.auto_ml.model.estimator

        # robust feature tracking
        names = None
        if hasattr(self.model_, "get_booster"):
            booster = self.model_.get_booster()
            if booster is not None:
                names = booster.feature_names
        if names is None and hasattr(self.model_, "feature_names_in_"):
            names = list(self.model_.feature_names_in_)
        # fallback: keep incoming order
        self.training_features_ = names if names is not None else list(X_df.columns)
        self.n_expected_ = getattr(self.model_, "n_features_in_", len(self.training_features_))
        return self


    def predict(self, X):
        X_df = pd.DataFrame(X)
        # ensure all expected cols
        for col in self.training_features_:
            if col not in X_df.columns:
                X_df[col] = 0.0
        X_df = X_df[self.training_features_]
        # align width to what the estimator finally used (after dropping constants)
        if X_df.shape[1] > self.n_expected_:
            X_df = X_df.iloc[:, : self.n_expected_]
        elif X_df.shape[1] < self.n_expected_:
            pad = np.zeros((len(X_df), self.n_expected_ - X_df.shape[1]))
            X_df = pd.concat([X_df, pd.DataFrame(pad, index=X_df.index)], axis=1)
        return self.model_.predict(X_df)

# 4) Robust AutoML builder 

def make_robust_automl(time_budget=120, metric="rmse"):
    """FixedSchema ➜ Imputer ➜ Scaler ➜ FLAML"""
    return make_pipeline(
        FixedSchema(fill_value=0.0, as_ndarray=False),
        SimpleImputer(strategy="median"),
        StandardScaler(),
        FlamlRegressorDoubleML(time_budget=time_budget, metric=metric),
    )

# 5) Data load & preprocessing
logger.info("Loading data …")
df = df, parse_dates=["Date"])
df["Amount spent (USD)"].ffill(inplace=True)

# Past‑only ad‑stock → no leakage
alpha, lags = 0.9, 7
w = alpha ** np.arange(lags)
df["adstock_spend"] = (
    df["Amount spent (USD)"]
    .rolling(window=lags, min_periods=1)
    .apply(lambda x: np.dot(x, w[len(w) - len(x):]), raw=True)
)

df = df.sort_values("Date").dropna(subset=["gmv", "adstock_spend"]).reset_index(drop=True)
logger.info("Data ready: %d obs", len(df))

# 3) Feature engineering & DoubleML data


T0_FE = time.time()
featurizer = PureTemporalFeatures("Date")
df_feat = featurizer.fit_transform(df)

business_cols = ["10_yr_ir", "dj_close_open_ratio", "App DAUs (iOS+Android)", "posts"]

temp_feats = [
    "month_sin","month_cos","dow_sin","dow_cos","doy_sin","doy_cos","is_weekend","is_month_start",
    "is_month_end","is_quarter_start","is_quarter_end","is_holiday_season","is_black_friday_week",
    "is_nordstrom_anniversary_sale","is_sephora_spring_sale","is_walmart_july_event","is_target_july_event",
    "is_amazon_prime_days_window","time_trend","time_trend_sq",
]

all_x_cols = business_cols + temp_feats

# build modelling frame -------------------------------------------------
df_model = df_feat.copy()
df_model["gmv"]            = df["gmv"]
df_model["adstock_spend"]  = df["adstock_spend"]

df_model.dropna(subset=["gmv","adstock_spend"], inplace=True)
df_model.reset_index(drop=True, inplace=True)

dml_data = dml.DoubleMLData(
    df_model,
    y_col="gmv",
    d_cols=["adstock_spend"],
    x_cols=all_x_cols,
)
logger.info("Feature engineering & DMLData creation took %.2fs — %d obs × %d covs",
            time.time()-T0_FE, dml_data.n_obs, len(dml_data.x_cols))


# 4) Manual, time-aware CV to obtain external predictions


logger.info("Generating external predictions with TimeSeriesSplit")
t0_manual_cv = time.time()

tscv = TimeSeriesSplit(n_splits=N_FOLDS_TS)
robust_l = make_robust_automl(time_budget=60, metric="rmse")
robust_m = make_robust_automl(time_budget=60, metric="rmse")

preds_l_ext = np.full((dml_data.n_obs, 1), np.nan)
preds_m_ext = np.full((dml_data.n_obs, 1), np.nan)

for k, (tri, tei) in enumerate(tscv.split(dml_data.x), 1):
    logger.info("Fold %d: train %d–%d | test %d–%d", k, tri[0], tri[-1], tei[0], tei[-1])

    # outcome model ------------------------------------------------------
    robust_l.fit(dml_data.x[tri, :], dml_data.y[tri])
    preds_l_ext[tei, 0] = robust_l.predict(dml_data.x[tei, :])

    # treatment model ----------------------------------------------------
    robust_m.fit(dml_data.x[tri, :], dml_data.d[tri])
    preds_m_ext[tei, 0] = robust_m.predict(dml_data.x[tei, :])

logger.info("CV finished in %.2fs", time.time() - t0_manual_cv)
assert not np.isnan(preds_l_ext).any()
assert not np.isnan(preds_m_ext).any()

external_predictions_dict = {
    "adstock_spend": {"ml_l": preds_l_ext, "ml_m": preds_m_ext}
}


# 5) DoubleMLPLR with external predictions


t0_dml = time.time()
logger.info("Fitting DoubleMLPLR with external predictions")
dml_plr = dml.DoubleMLPLR(
    dml_data,
    ml_l=None,               # not used because we pass external preds
    ml_m=None,
    n_folds=N_FOLDS_TS,
    n_rep=N_REP_MANUAL,
    score="partialling out",
    draw_sample_splitting=False,
)

res = dml_plr.fit(external_predictions=external_predictions_dict, store_models=False)
logger.info("DoubleMLPLR completed in %.2fs", time.time() - t0_dml)

print("\nCausal estimate (theta):", res.coef)
print("Std. error:", res.se)

logger.info("Total script time: %.2fs", time.time() - T0_TOTAL)

#7) Results & Diagnostics (with External Predictions) 
logger.info("Causal effect estimate (from external predictions):")
logger.info(f"\n{res_ext.summary}")

print("\n--- External Prediction Run (Time-Series CV) ---")
try:
    print("Fit with external predictions completed successfully.")
    print("Coefficients (theta):", res_ext.coef)
    print("Standard errors:", res_ext.se)
    print(dml_plr)

    print("\nEvaluating learners (based on manual OOS predictions from TimeSeriesSplit):")
    if oos_rmse_l_scores:
        mean_oos_rmse_l = np.mean(oos_rmse_l_scores)
        std_oos_rmse_l = np.std(oos_rmse_l_scores)
        print(f"Mean OOS RMSE (ml_l) across TS folds: {mean_oos_rmse_l:.4f}")
        print(f"Std OOS RMSE (ml_l) across TS folds: {std_oos_rmse_l:.4f}")
        # print(f"  All ml_l RMSEs: {[f'{s:.4f}' for s in oos_rmse_l_scores]}")
    if oos_rmse_m_scores:
        mean_oos_rmse_m = np.mean(oos_rmse_m_scores)
        std_oos_rmse_m = np.std(oos_rmse_m_scores)
        print(f"Mean OOS RMSE (ml_m) across TS folds: {mean_oos_rmse_m:.4f}")
        print(f"Std OOS RMSE (ml_m) across TS folds: {std_oos_rmse_m:.4f}")
        # print(f"  All ml_m RMSEs: {[f'{s:.4f}' for s in oos_rmse_m_scores]}")

except Exception as e:
    print(f" Post-fit error (external predictions): {e}")

logger.info(f"Total script time: {time.time() - t0_total:.2f}s")

[SvenKlaassen]

Could you adapt your example such that you use a simple generated dataset (so I can test it)?

[paul-jdfagan]

Hey Sven,

The only fully time-aware solution I can make work is to decouple the approach: nuisance models (ml_l and ml_m), internally tuned by FLAML using its split_type="time" cross-validation, trained once on an initial historical data segment to generate a single set of leak-safe "external predictions" for the entire dataset. These predictions are passed to DoubleML, where the n_rep parameter was adjusted to 1 to match the single prediction set, DoubleML performs its cross-fitting on the residuals derived from these pre-computed, time-aware predictions, hopefully! thus ensuring both temporal integrity and valid causal estimation.

===== Results =====
θ = 2.14
SE = 0.20
True θ = 2.50

"""
Pre-quential DoubleML demo (leak-safe, no cross-fitting)
Truth θ = 2.50
"""

# ------------------------------------------------------------------
# Imports & synthetic data 
# ------------------------------------------------------------------
import numpy as np, pandas as pd, random, logging
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.dummy import DummyRegressor
from flaml import AutoML
import doubleml as dml

np.random.seed(42)
random.seed(42)
logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

# 1. Synthetic generator 
N_DAYS = 1_000
dates  = pd.date_range("2021-01-01", periods=N_DAYS, freq="D")
alpha, LAGS, true_theta = 0.9, 7, 2.50

df = pd.DataFrame({"Date": dates})
df["10_yr_ir"]            = np.random.normal(2, 0.2,  N_DAYS)
df["dj_close_open_ratio"] = np.random.normal(1, 0.02, N_DAYS)
df["App DAUs (iOS+Android)"] = np.random.normal(50_000, 5_000, N_DAYS)
df["posts"]               = np.random.poisson(80, N_DAYS)
df["Amount spent (USD)"]  = np.random.gamma(5, 200,  N_DAYS)

w = alpha ** np.arange(LAGS)
df["adstock_spend"] = (
    df["Amount spent (USD)"]
    .rolling(LAGS, 1)
    .apply(lambda x: np.dot(x, w[len(w)-len(x):]), raw=True)
)

baseline    = 100_000
seasonality = 10_000*np.sin(2*np.pi*df["Date"].dt.dayofyear/365.25)
noise       = np.random.normal(0, 5_000, N_DAYS)
df["gmv"]   = baseline + true_theta*df["adstock_spend"] + seasonality + noise
logger.info("Synthetic data ready: %d rows", len(df))

# 2. Feature builder 
class PureTemporalFeatures(BaseEstimator, TransformerMixin):
    def __init__(self, date_col="Date"): self.date_col = date_col
    def fit(self, X, y=None): return self
    def transform(self, X):
        X = X.copy()
        X[self.date_col] = pd.to_datetime(X[self.date_col])
        X.sort_values(self.date_col, inplace=True)
        X["month_sin"] = np.sin(2*np.pi*X[self.date_col].dt.month/12)
        X["month_cos"] = np.cos(2*np.pi*X[self.date_col].dt.month/12)
        X["dow_sin"]   = np.sin(2*np.pi*X[self.date_col].dt.weekday/7)
        X["dow_cos"]   = np.cos(2*np.pi*X[self.date_col].dt.weekday/7)
        X["doy_sin"]   = np.sin(2*np.pi*X[self.date_col].dt.dayofyear/365.25)
        X["doy_cos"]   = np.cos(2*np.pi*X[self.date_col].dt.dayofyear/365.25)
        X["time_trend"]    = np.arange(len(X))
        X["time_trend_sq"] = X["time_trend"]**2
        return X.drop(columns=[self.date_col])

# 3. FixedSchema & FLAML wrapper 
class FixedSchema(BaseEstimator, TransformerMixin):
    def __init__(self, fill_value=0.0): self.fill_value = fill_value
    def fit(self, X, y=None):
        self.columns_ = list(X.columns)
        return self
    def transform(self, X):
        Xf = X.copy()
        for c in set(self.columns_) - set(Xf.columns):
            Xf[c] = self.fill_value
        return Xf[self.columns_]

class FlamlRegressor(BaseEstimator):
    _estimator_type = "regressor"
    def __init__(self, time_budget=20):
        self.time_budget = time_budget
        self.auto_ml = AutoML()
    def fit(self, X, y):
        # Ensure column names are strings for FLAML
        X_df = pd.DataFrame(X)
        X_df.columns = [str(col) for col in X_df.columns]
        self.auto_ml.fit(
            X_train=X_df, y_train=y, task="regression",
            time_budget=self.time_budget, split_type="time", # Split_type 'time' is for FLAML's internal CV
            metric="rmse", verbose=0
        )
        self.model_ = self.auto_ml.model.estimator
        # Store stringified column names for prediction
        self.training_features_ = list(X_df.columns)
        self.n_expected_ = getattr(self.model_, "n_features_in_", len(self.training_features_))
        return self
    def predict(self, X):
        X_df = pd.DataFrame(X)
        X_df.columns = [str(col) for col in X_df.columns] # Ensure prediction columns are also strings

        # Reindex to match training features, adding missing columns as 0.0
        X_df = X_df.reindex(columns=self.training_features_).fillna(0.0)

        # Ensure numeric width matches what the estimator finally expects
        current_n_features = X_df.shape[1]
        if current_n_features > self.n_expected_:
            X_df = X_df.iloc[:, :self.n_expected_]
        elif current_n_features < self.n_expected_:
            pad_width = self.n_expected_ - current_n_features
            pad_df = pd.DataFrame(np.zeros((len(X_df), pad_width)), index=X_df.index,
                                    columns=[f"pad_{i}" for i in range(pad_width)])
            X_df = pd.concat([X_df, pad_df], axis=1)

        return self.model_.predict(X_df.to_numpy())

def make_automl():
    return make_pipeline(
        FixedSchema(fill_value=0.0), # FixedSchema ensures column consistency
        SimpleImputer(strategy="median"),
        StandardScaler(),
        FlamlRegressor(time_budget=20)
    )

# 4. Model frame & DoubleML container 
business = ["10_yr_ir","dj_close_open_ratio","App DAUs (iOS+Android)","posts"]
temps    = ["month_sin","month_cos","dow_sin","dow_cos","doy_sin","doy_cos",
            "time_trend","time_trend_sq"]
x_cols   = business + temps

feat  = PureTemporalFeatures("Date")
df_x  = feat.fit_transform(df)
df_x["gmv"]           = df["gmv"]
df_x["adstock_spend"] = df["adstock_spend"]

dml_data = dml.DoubleMLData(
    df_x, y_col="gmv", d_cols=["adstock_spend"], x_cols=x_cols
)

# 5. Pre-quential split & nuisance models
# Define the split for pre-quential prediction
# This defines the "training cutoff" for external predictions TRAIN_END = 750
# The remaining part of the data will be used for testing. For a true pre-quential setup, you'd typically have multiple training/test splits, but for this simplified example, we'll use a single split for external predictions that covers the entire dataset.

# Train your learners on the initial training segment
ml_l = make_automl()
ml_m = make_automl()

# Convert to DataFrame for FixedSchema and FLAML's internal use (especially for column names)
X_train_df = pd.DataFrame(dml_data.x[:TRAIN_END, :], columns=x_cols)

ml_l.fit(X_train_df, dml_data.y[:TRAIN_END])
ml_m.fit(X_train_df, dml_data.d[:TRAIN_END])

# Generate predictions for the *entire* dataset using the models trained
# on the initial segment.
# Convert the entire dml_data.x to DataFrame for prediction consistency
X_full_df = pd.DataFrame(dml_data.x[:, :], columns=x_cols)

preds_l = ml_l.predict(X_full_df)[:, None]
preds_m = ml_m.predict(X_full_df)[:, None]

external = {"adstock_spend": {"ml_l": preds_l, "ml_m": preds_m}}

# 6. DoubleML PLR 
# For DoubleML, when using external_predictions,  use its own cross-fitting scheme (default K-fold) to aggregate the nuisance estimates.

N_FOLDS_DML = 5 # DoubleML's default number of folds for internal cross-fitting

plr = dml.DoubleMLPLR(
    dml_data,
    ml_l=DummyRegressor(),
    ml_m=DummyRegressor(),
    n_folds=N_FOLDS_DML,
    n_rep=1, # Since you're providing predictions for all data points, 1 repetition is usually fine.
    score="partialling out",
    draw_sample_splitting=True # Allow DoubleML to generate standard cross-fitting splits
                               # to aggregate the externally provided predictions.
)

res = plr.fit(external_predictions=external, store_models=False)

print("\n===== Results =====")
print(f"θ  = {res.coef[0]:.2f}")
print(f"SE = {res.se[0]:.2f}")
print(f"True θ = {true_theta:.2f}")

[SvenKlaassen]

@paul-jdfagan, thanks for the example.

Working with pre-defined sample splits seems not possible in a time-series setting as the package currently expects scores for all observations (which is not possible for the first training period, even if the cv-split respects the time structure).

That looks already fine to me.
The only extension to your example would be a rolling/expanding window type prediction setting. This would respect the time-structure but allow for a larger sample size of the predictions.

[OliverSchacht]

Hi @paul-jdfagan thanks for the example.

I think the suggested solution by @SvenKlaassen works great.

If your use case requires more sophisticated splitting, i would recommend you to work with a custimized splitter as proposed in this discussion

As you see in the code, you can hand over a custom TimeSeriesSplit()-object to flaml in order to get reproducable splits. This is what flaml does internally, too, when you call it as above. You can then use the splitter to generate out-of-sample external predictions for DoubleML. You can also fine tune more.

Your code could look something like below. An advantage would be that this is more data efficient potentially as you only throw out the first $n/k$ samples in $k$-fold CV. On the other side, it uses a different CV scheme than default DoubleML. You might want to read up the caveats in the double machine learning literature.

import numpy as np
import pandas as pd
from sklearn.model_selection import TimeSeriesSplit
from sklearn.metrics import root_mean_squared_error

from flaml import AutoML

# create dummy time series
n = 200
date_idx = pd.date_range("2021-01-01", periods=n, freq="D")
X = pd.DataFrame({
        "trend": np.arange(n),
        "sin":   np.sin(np.arange(n)/12),
        "cos":   np.cos(np.arange(n)/12)
    }, index=date_idx)
rng = np.random.RandomState(42)
y = 0.3 * X["trend"] + X["sin"] + rng.normal(scale=0.5, size=n)

# define custom splitter (equivalent to setting `split_type="time"` in flaml)
my_splitter = TimeSeriesSplit(n_splits=5) # you can do customizations here

# define AutoML
automl = AutoML()
automl_settings = dict(
    task="regression",
    metric="rmse",           
    time_budget=60,            
    split_type=my_splitter,
    verbose=0,
)
automl.fit(X_train=X, y_train=y, **automl_settings)

# Create predictions with same splits
cv_preds = np.full_like(y, fill_value=np.nan, dtype=float)

best_estimator = automl.model.estimator

for train_idx, test_idx in my_splitter.split(X):
    best_estimator.fit(X.iloc[train_idx], y.iloc[train_idx])
    cv_preds[test_idx] = best_estimator.predict(X.iloc[test_idx])

mask = ~np.isnan(cv_preds)
cv_loss = root_mean_squared_error(y[mask], cv_preds[mask])

# mask first data that is never predicted before creating DoubleMLData object
X_for_doubleml = X.iloc[mask]
y_for_doubleml = y[mask]

Best, Oliver

[paul-jdfagan]

Hi @SvenKlaassen and @OliverSchacht,
Thanks so much for taking the time to help me work through my problem and offer solutions. I really appreciate it. Lots to think about.

Apologies if i'm missing something obvious to you both. To me, the pipeline can be fully time-aware (via various routes), and the code can run, right up until the final step of integrating sample_splitting into DoubleML (code 1) then I keep getting the dreaded "ValueError: Invalid partition provided. Tuples provided that don't form a partition".

I think I might I have to accept a middle ground in the tradeoff between being fully time-aware or fully theoretically sound for orthogonalisation - thanks Oliver. The Code 2 flow is probably one I can live with. Explicitly enforces time-aware splits in most steps, but ultimately runs with CV in DML.

CODE 1
# ─── 9) Prepare sample splitting aligned with FLAML ─────────────────────────
logger.info("Preparing sample splits for DoubleML to align with FLAML")

# Recreate the TimeSeriesSplit on the masked data
masked_splitter = TimeSeriesSplit(n_splits=n_folds)
smpls = [(train_idx, test_idx) for train_idx, test_idx in masked_splitter.split(X_for_doubleml)]

# ─── 10) Set up DoubleML with tuned learners and injected sample splitting ──
logger.info("Setting up DoubleML with aligned sample splitting")

ml_l = make_pipeline(
    SimpleImputer(strategy="median"),
    StandardScaler(),
    FlamlRegressorDoubleML(time=60, metric="rmse")
)
ml_m = make_pipeline(
    SimpleImputer(strategy="median"),
    StandardScaler(),
    FlamlRegressorDoubleML(time=60, metric="rmse")
)

dml_plr = dml.DoubleMLPLR(
    dml_data,
    ml_l=ml_l,
    ml_m=ml_m,
    n_folds=n_folds,
    n_rep=1,
    score="partialling out",
    draw_sample_splitting=False
)

#  Inject your custom sample splits from FLAML
dml_plr.set_sample_splitting(smpls)

CODE 2

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.model_selection import TimeSeriesSplit
from sklearn.metrics import root_mean_squared_error
from flaml import AutoML
import doubleml as dml
import logging
import time

# Configure Logging 
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 1) Pure temporal transformer 
class PureTemporalFeatures(BaseEstimator, TransformerMixin):
    """
    Add only date-based, leak-safe features.
    """
    def __init__(self, date_col="Date"):
        self.date_col = date_col

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        t0 = time.time()
        logger.info("Starting PureTemporalFeatures.transform")
        df = X.copy()
        df[self.date_col] = pd.to_datetime(df[self.date_col])
        df = df.sort_values(self.date_col).reset_index(drop=True)

        # Basic cyclical encodings
        df["month_sin"] = np.sin(2*np.pi * df[self.date_col].dt.month/12)
        df["month_cos"] = np.cos(2*np.pi * df[self.date_col].dt.month/12)
        df["dow_sin"]   = np.sin(2*np.pi * df[self.date_col].dt.weekday/7)
        df["dow_cos"]   = np.cos(2*np.pi * df[self.date_col].dt.weekday/7)
        df["doy_sin"]   = np.sin(2*np.pi * df[self.date_col].dt.dayofyear/365)
        df["doy_cos"]   = np.cos(2*np.pi * df[self.date_col].dt.dayofyear/365)

        # calendar flags
        df["is_weekend"]       = (df[self.date_col].dt.weekday >= 5).astype(int)
        df["is_month_start"]   = (df[self.date_col].dt.day <= 7).astype(int)
        df["is_month_end"]     = (df[self.date_col].dt.day >= 24).astype(int)
        df["is_quarter_start"] = df[self.date_col].dt.month.isin([1,4,7,10]).astype(int)
        df["is_quarter_end"]   = df[self.date_col].dt.month.isin([3,6,9,12]).astype(int)

        # shopping-season flags
        df["is_holiday_season"]           = (
            (df[self.date_col].dt.month==12) |
            ((df[self.date_col].dt.month==11)&(df[self.date_col].dt.day>=15))
        ).astype(int)
        df["is_black_friday_week"]        = (
            (df[self.date_col].dt.month==11)&(df[self.date_col].dt.day>=20)
        ).astype(int)
        df["is_nordstrom_anniversary_sale"] = (
            ((df[self.date_col].dt.month==7)&(df[self.date_col].dt.day>=15)) |
            ((df[self.date_col].dt.month==8)&(df[self.date_col].dt.day<=10))
        ).astype(int)
        df["is_sephora_spring_sale"]      = (
            (df[self.date_col].dt.month==4)&
            df[self.date_col].dt.day.between(1,20)
        ).astype(int)
        df["is_walmart_july_event"]       = (
            (df[self.date_col].dt.month==7)&
            df[self.date_col].dt.day.between(10,20)
        ).astype(int)
        df["is_target_july_event"]        = df["is_walmart_july_event"]
        df["is_amazon_prime_days_window"] = (
            (df[self.date_col].dt.month==7)&
            df[self.date_col].dt.day.between(9,18)
        ).astype(int)

        # simple trend
        df["time_trend"]    = np.arange(len(df))
        df["time_trend_sq"] = df["time_trend"] ** 2

        elapsed = time.time() - t0
        logger.info(f"PureTemporalFeatures.transform completed in {elapsed:.2f}s, added {len(df.columns)-1} features")
        return df.drop(columns=[self.date_col])


#  2) FLAML wrapper for DoubleML
class FlamlRegressorDoubleML:
    _estimator_type = "regressor"
    def __init__(self, time=120, estimator_list=None, metric="rmse", **kwargs):
        self.auto_ml = AutoML(**kwargs)
        self.time = time
        self.estimator_list = estimator_list or ["lgbm","xgboost","histgb","xgb_limitdepth","catboost"]
        self.metric = metric
        self.best_loss_ = None
        self.tuned_model = None

    def set_params(self, **params):
        self.auto_ml.set_params(**params)
        return self

    def get_params(self, deep=True):
        p = self.auto_ml.get_params(deep)
        p.update(time=self.time, estimator_list=self.estimator_list, metric=self.metric)
        return p

    def fit(self, X, y):
        t0 = time.time()
        logger.info(f"FLAML.fit: starting on {X.shape[0]} rows, budget={self.time}s")
        self.auto_ml.fit(
            X, y,
            task="regression",
            time_budget=self.time,
            estimator_list=self.estimator_list,
            metric=self.metric,
            split_type="time",
            verbose=False
        )
        self.tuned_model = self.auto_ml.model.estimator
        try:
            self.best_loss_ = self.auto_ml.best_loss
        except:
            self.best_loss_ = getattr(self.auto_ml._state, "best_loss", None)
        elapsed = time.time() - t0
        logger.info(f"FLAML.fit completed in {elapsed:.2f}s, best_loss_={self.best_loss_}")
        return self

    def predict(self, X):
        return self.tuned_model.predict(X)


# 3) Load & preprocess 
t0_total = time.time()
logger.info("Starting data load & preprocessing")
df = df
df["Amount spent (USD)"].fillna(method='ffill', inplace=True)

def geometric_adstock(s, alpha=0.9, lags=7):
    w = alpha ** np.arange(lags)
    return s.rolling(lags, min_periods=1) \
            .apply(lambda x: np.dot(x[::-1], w[:len(x)]), raw=True)

df["adstock_spend"] = geometric_adstock(df["Amount spent (USD)"])
df.dropna(subset=["gmv","adstock_spend"], inplace=True)
df.reset_index(drop=True, inplace=True)
logger.info(f"Data loaded & preprocessed in {time.time() - t0_total:.2f}s — {len(df)} obs")

# 4) Feature engineering 
t0 = time.time()
logger.info("Applying PureTemporalFeatures")
df_temp = PureTemporalFeatures(date_col="Date").fit_transform(df)

business = ["10_yr_ir","dj_close_open_ratio","App DAUs (iOS+Android)","posts"]
temp_feats = [
    "month_sin","month_cos","dow_sin","dow_cos","doy_sin","doy_cos",
    "is_weekend","is_month_start","is_month_end",
    "is_quarter_start","is_quarter_end","is_holiday_season","is_black_friday_week",
    "is_nordstrom_anniversary_sale","is_sephora_spring_sale",
    "is_walmart_july_event","is_target_july_event","is_amazon_prime_days_window",
    "time_trend","time_trend_sq"
]
all_x = business + temp_feats

df_model = pd.concat([
    df_temp[temp_feats],
    df[business],
    df[["gmv","adstock_spend"]]
], axis=1).loc[:, lambda d: ~d.columns.duplicated()]

logger.info(f"Feature engineering took {time.time() - t0:.2f}s — {len(df_model)} obs × {len(all_x)} covs")

# 5) Custom TimeSeriesSplit approach 
n_folds = 5
my_splitter = TimeSeriesSplit(n_splits=n_folds)

logger.info(f"Using TimeSeriesSplit with {n_folds} folds on {len(df_model)} observations")

# Log the fold structure for transparency
logger.info("Time series fold structure:")
for fold_idx, (train_idx, test_idx) in enumerate(my_splitter.split(df_model)):
    logger.info(f"  Fold {fold_idx + 1}: Train={len(train_idx)} obs (idx {train_idx[0]}-{train_idx[-1]}), "
                f"Test={len(test_idx)} obs (idx {test_idx[0]}-{test_idx[-1]})")

#6) Fit FLAML models with custom splitter 
t0 = time.time()
logger.info("Fitting FLAML models for nuisance functions")

# Prepare data for FLAML
X_all = df_model[all_x]
y_outcome = df_model["gmv"]  # outcome variable
d_treatment = df_model["adstock_spend"]  # treatment variable

# Fit ml_l (outcome model) using custom splitter
logger.info("Fitting ml_l (outcome model)")
automl_l = AutoML()
automl_l_settings = dict(
    task="regression",
    metric="rmse",
    time_budget=120,
    split_type=my_splitter,
    estimator_list=["lgbm","xgboost","histgb","xgb_limitdepth","catboost"],
    verbose=0,
)
automl_l.fit(X_train=X_all, y_train=y_outcome, **automl_l_settings)

# Fit ml_m (treatment model) using custom splitter  
logger.info("Fitting ml_m (treatment model)")
automl_m = AutoML()
automl_m_settings = dict(
    task="regression",
    metric="rmse",
    time_budget=120,
    split_type=my_splitter,
    estimator_list=["lgbm","xgboost","histgb","xgb_limitdepth","catboost"],
    verbose=0,
)
automl_m.fit(X_train=X_all, y_train=d_treatment, **automl_m_settings)

logger.info(f"FLAML models fitted in {time.time() - t0:.2f}s")

# 7) Generate cross-validated predictions 
t0 = time.time()
logger.info("Generating cross-validated predictions")

# Create CV predictions using the same splits
cv_preds_l = np.full_like(y_outcome, fill_value=np.nan, dtype=float)
cv_preds_m = np.full_like(d_treatment, fill_value=np.nan, dtype=float)

best_estimator_l = automl_l.model.estimator
best_estimator_m = automl_m.model.estimator

for fold_idx, (train_idx, test_idx) in enumerate(my_splitter.split(X_all)):
    logger.info(f"Processing fold {fold_idx + 1}/{n_folds}")
    
    # Fit and predict for outcome model (ml_l)
    best_estimator_l.fit(X_all.iloc[train_idx], y_outcome.iloc[train_idx])
    cv_preds_l[test_idx] = best_estimator_l.predict(X_all.iloc[test_idx])
    
    # Fit and predict for treatment model (ml_m)
    best_estimator_m.fit(X_all.iloc[train_idx], d_treatment.iloc[train_idx])
    cv_preds_m[test_idx] = best_estimator_m.predict(X_all.iloc[test_idx])

# Create mask for valid predictions (exclude first observations that were never predicted)
mask = ~np.isnan(cv_preds_l) & ~np.isnan(cv_preds_m)
logger.info(f"Valid predictions for {mask.sum()}/{len(mask)} observations")

# Calculate CV losses
cv_loss_l = root_mean_squared_error(y_outcome[mask], cv_preds_l[mask])
cv_loss_m = root_mean_squared_error(d_treatment[mask], cv_preds_m[mask])
logger.info(f"CV RMSE - ml_l: {cv_loss_l:.4f}, ml_m: {cv_loss_m:.4f}")

logger.info(f"Cross-validated predictions generated in {time.time() - t0:.2f}s")

# 8) Prepare data for DoubleML 
X_for_doubleml = X_all.iloc[mask]
y_for_doubleml = y_outcome.iloc[mask]
d_for_doubleml = d_treatment.iloc[mask]

df_for_doubleml = pd.concat([
    X_for_doubleml,
    pd.Series(y_for_doubleml, name="gmv", index=X_for_doubleml.index),
    pd.Series(d_for_doubleml, name="adstock_spend", index=X_for_doubleml.index)
], axis=1)

dml_data = dml.DoubleMLData(df_for_doubleml, y_col="gmv", d_cols="adstock_spend", x_cols=all_x)
logger.info(f"DoubleMLData created with {dml_data.n_obs} observations")

#  9) Set up DoubleML with regular learners 
t0 = time.time()
logger.info("Setting up DoubleML with standard learners")

# Use simpler learners since we've already done the hyperparameter tuning
ml_l = make_pipeline(
    SimpleImputer(strategy="median"),
    StandardScaler(),
    FlamlRegressorDoubleML(time=60, metric="rmse")
)
ml_m = make_pipeline(
    SimpleImputer(strategy="median"),
    StandardScaler(),
    FlamlRegressorDoubleML(time=60, metric="rmse")
)

# 10) Fit DoubleMLPLR
logger.info("Fitting DoubleMLPLR")
dml_plr = dml.DoubleMLPLR(
    dml_data,
    ml_l=ml_l,
    ml_m=ml_m,
    n_folds=5,
    n_rep=3,
    score="partialling out",
    draw_sample_splitting=True
)

res = dml_plr.fit(n_jobs_cv=-1, store_models=True)
logger.info(f"DoubleMLPLR fit completed in {time.time() - t0:.2f}s")

# 11) Results & diagnostics 
logger.info("Causal effect estimate:")
logger.info(f"\n{res.summary}")
print(dml_plr)

try:
    print("Fit completed successfully. Coefficients:", res.coef)

    # ─── Evaluate Learners ───────────────────────────────────────────────
    print("\nEvaluating learners...")
    try:
        evals = dml_plr.evaluate_learners(learners=['ml_l', 'ml_m'])
        oos_rmse_l = np.mean([v[0] for v in evals['ml_l']])
        oos_rmse_m = np.mean([v[0] for v in evals['ml_m']])
        print(f"OOS RMSE (ml_l): {oos_rmse_l:.4f}")
        print(f"OOS RMSE (ml_m): {oos_rmse_m:.4f}")
        
        # Compare with our custom CV results
        print(f"\nComparison with custom TimeSeriesSplit CV:")
        print(f"Custom CV RMSE (ml_l): {cv_loss_l:.4f}")
        print(f"Custom CV RMSE (ml_m): {cv_loss_m:.4f}")
        
        # Detailed evaluation metrics
        print("\nDetailed learner evaluation:")
        for learner_name in ['ml_l', 'ml_m']:
            if learner_name in evals:
                scores = evals[learner_name]
                rmse_scores = [score[0] for score in scores]
                print(f"{learner_name}:")
                print(f"  RMSE scores across folds: {[f'{score:.4f}' for score in rmse_scores]}")
                print(f"  Mean RMSE: {np.mean(rmse_scores):.4f}")
                print(f"  Std RMSE: {np.std(rmse_scores):.4f}")
                cv_value = np.std(rmse_scores)/np.mean(rmse_scores) if np.mean(rmse_scores) != 0 else np.inf
                print(f"  CV (std/mean): {cv_value:.4f}")
    except Exception as e:
        print(f"Error in evaluating learners: {e}")

    #  Additional diagnostics 
    print(f"\nData Processing Summary:")
    print(f"  Original observations: {len(df_model)}")
    print(f"  Valid CV predictions: {mask.sum()} ({100*mask.sum()/len(df_model):.1f}%)")
    print(f"  Used for DoubleML: {dml_data.n_obs}")
    print(f"  Time series folds: {n_folds}")
    print(f"  Data efficiency: ~{100 * (n_folds-1) / n_folds:.1f}% of data used for training")

except Exception as e:
    print("Post-fit error:", e)

logger.info(f"Total script time: {time.time() - t0_total:.2f}s")

[SvenKlaassen]

@paul-jdfagan,

The reason why the first example is not working, is due to internal validation checks of the package. For each sample split, we validate if the number of predictions equals the length of the outcome vector etc.
This could be adapated in the future but requires some larger changes in the package.

I think your second example is fine, but if you instead would like to use the cross-fitted predictions from flaml directly, you can just "reduce" your data for DoulbeML to the relevant size. The first block of $n/k$ observations has np.nan predictions if the time structure is respected.
Therefore, we can drop them and just use the rest for the DoubleMLData object.
Starting from your step $8$, you could instead directly use the cross-fitted predictions form flaml:

# 8) Prepare data for DoubleML 
y_for_doubleml = y_outcome.iloc[mask]
d_for_doubleml = d_treatment.iloc[mask]

df_for_doubleml = pd.concat([
    pd.Series(y_for_doubleml, name="gmv", index=y_for_doubleml.index),
    pd.Series(d_for_doubleml, name="adstock_spend", index=y_for_doubleml.index),
   pd.Series(np.zeros_like(y_for_doubleml), name="dummy_x", index=y_for_doubleml.index),
], axis=1)

dml_data = dml.DoubleMLData(df_for_doubleml, y_col="gmv", d_cols="adstock_spend", x_cols=["dummy_x"])
logger.info(f"DoubleMLData created with {dml_data.n_obs} observations")

# 9) Prepare Dummy Learners & external predictions

ml_l_dummy = dml.utils.DMLDummyRegressor()
ml_m_dummy = dml.utils.DMLDummyRegressor()

pred_dict = {"adstock_spend": {
    "ml_l": cv_preds_l[mask].reshape(-1, 1),
    "ml_m": cv_preds_m[mask].reshape(-1, 1),
    }
}

# 10) Fit DoubleMLPLR
t0 = time.time()
logger.info("Fitting DoubleMLPLR")
dml_plr = dml.DoubleMLPLR(
    dml_data,
    ml_l=ml_l_dummy,
    ml_m=ml_m_dummy,
    n_folds=5,
    n_rep=1,
    score="partialling out",
)

res = dml_plr.fit(external_predictions=pred_dict)
logger.info(f"DoubleMLPLR fit completed in {time.time() - t0:.2f}s")

This way you keep the time, relevant splitted predictions.