From 1006540fd3ed59048ebe0c3887021aa0793d792f Mon Sep 17 00:00:00 2001 From: ***** <721466+soodoku@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:20:51 -0800 Subject: [PATCH] test: validate LPLR tuning on its optimization folds --- .../plm/tests/test_lplr_tune_ml_models.py | 114 +++++++++--------- 1 file changed, 54 insertions(+), 60 deletions(-) diff --git a/doubleml/plm/tests/test_lplr_tune_ml_models.py b/doubleml/plm/tests/test_lplr_tune_ml_models.py index eb44b163..25f9957b 100644 --- a/doubleml/plm/tests/test_lplr_tune_ml_models.py +++ b/doubleml/plm/tests/test_lplr_tune_ml_models.py @@ -1,15 +1,15 @@ +"""Validate LPLR tuning against its cross-validation objective and applied parameters.""" + import numpy as np +import optuna import pytest +from sklearn.base import clone +from sklearn.model_selection import KFold, cross_val_score from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor import doubleml as dml from doubleml.plm.datasets import make_lplr_LZZ2020 -from doubleml.tests._utils_tune_optuna import ( - _SAMPLER_CASES, - _assert_tree_params, - _basic_optuna_settings, - _small_tree_params, -) +from doubleml.tests._utils_tune_optuna import _assert_tree_params, _basic_optuna_settings, _small_tree_params @pytest.fixture(scope="module", params=["nuisance_space", "instrument"]) @@ -17,77 +17,71 @@ def score(request): return request.param -@pytest.fixture( - scope="module", - params=[DecisionTreeRegressor(random_state=567), None], -) +@pytest.fixture(scope="module", params=[DecisionTreeRegressor(random_state=567), None]) def ml_a(request): return request.param @pytest.mark.ci -@pytest.mark.parametrize("sampler_name,optuna_sampler", _SAMPLER_CASES, ids=[case[0] for case in _SAMPLER_CASES]) -def test_doubleml_lplr_optuna_tune(sampler_name, optuna_sampler, score, ml_a): +@pytest.mark.parametrize("sampler_class", [optuna.samplers.RandomSampler, optuna.samplers.TPESampler], ids=["random", "tpe"]) +def test_doubleml_lplr_optuna_tune(sampler_class, score, ml_a): + """Tuning improves its own CV objective and applies the selected parameters.""" np.random.seed(3141) - alpha = 0.5 - dml_data = make_lplr_LZZ2020(n_obs=500, dim_x=15, alpha=alpha) - - ml_M = DecisionTreeClassifier(random_state=123) - ml_t = DecisionTreeRegressor(random_state=234) - ml_m = DecisionTreeRegressor(random_state=456) - + dml_data = make_lplr_LZZ2020(n_obs=500, dim_x=15, alpha=0.5) dml_lplr = dml.DoubleMLLPLR( dml_data, - ml_M=ml_M, - ml_t=ml_t, - ml_m=ml_m, + ml_M=DecisionTreeClassifier(random_state=123), + ml_t=DecisionTreeRegressor(random_state=234), + ml_m=DecisionTreeRegressor(random_state=456), ml_a=ml_a, n_folds=2, n_folds_inner=2, score=score, ) - dml_lplr.fit() - untuned_score = dml_lplr.evaluate_learners() - - optuna_params = { - "ml_M": _small_tree_params, - "ml_m": _small_tree_params, - "ml_t": _small_tree_params, - "ml_a": _small_tree_params, + cv = KFold(n_splits=5, shuffle=True, random_state=42) + x, y, d = dml_data.x, dml_data.y, dml_data.d + mask_m = y == 0 if score == "nuisance_space" else np.ones(len(y), dtype=bool) + tuning_data = { + "ml_M": (np.column_stack((d, x)), y, "neg_log_loss"), + "ml_m": (x[mask_m], d[mask_m], "neg_root_mean_squared_error"), + "ml_a": (x, d, "neg_root_mean_squared_error"), } - + baseline_scores = { + name: cross_val_score(clone(dml_lplr.learner[name]), features, target, cv=cv, scoring=scorer).mean() + for name, (features, target, scorer) in tuning_data.items() + } + optuna_params = {name: _small_tree_params for name in ["ml_M", "ml_m", "ml_a", "ml_t"]} tune_res = dml_lplr.tune_ml_models( ml_param_space=optuna_params, - optuna_settings=_basic_optuna_settings({"sampler": optuna_sampler, "n_trials": 5}), + cv=cv, + optuna_settings=_basic_optuna_settings({"sampler": sampler_class(seed=3141), "n_trials": 5}), return_tune_res=True, ) - dml_lplr.fit() - tuned_score = dml_lplr.evaluate_learners() - - tuned_params_M = tune_res[0]["ml_M"].best_params - tuned_params_t = tune_res[0]["ml_t"].best_params - tuned_params_m = tune_res[0]["ml_m"].best_params - tuned_params_a = tune_res[0]["ml_a"].best_params + assert len(tune_res) == 1 + assert set(tune_res[0]) == set(optuna_params) + for name, result in tune_res[0].items(): + assert result.tuned + _assert_tree_params(result.best_params) + assert dml_lplr.get_params(name)["d"] == [[result.best_params] * dml_lplr.n_folds] + for key, value in result.best_params.items(): + assert result.best_estimator.get_params()[key] == value + + for name, (features, target, scorer) in tuning_data.items(): + result = tune_res[0][name] + verified_score = cross_val_score(clone(result.best_estimator), features, target, cv=cv, scoring=scorer).mean() + assert result.scoring_method == scorer + assert result.best_score == pytest.approx(verified_score) + assert verified_score > baseline_scores[name] - _assert_tree_params(tuned_params_M) - _assert_tree_params(tuned_params_t) - _assert_tree_params(tuned_params_m) - _assert_tree_params(tuned_params_a) - - # ensure results contain optuna objects and best params - assert isinstance(tune_res[0], dict) - assert set(tune_res[0].keys()) == {"ml_M", "ml_m", "ml_t", "ml_a"} - assert hasattr(tune_res[0]["ml_M"], "best_params") - assert tune_res[0]["ml_M"].best_params["max_depth"] == tuned_params_M["max_depth"] - assert hasattr(tune_res[0]["ml_t"], "best_params") - assert tune_res[0]["ml_t"].best_params["max_depth"] == tuned_params_t["max_depth"] - assert hasattr(tune_res[0]["ml_m"], "best_params") - assert tune_res[0]["ml_m"].best_params["max_depth"] == tuned_params_m["max_depth"] - assert hasattr(tune_res[0]["ml_a"], "best_params") - assert tune_res[0]["ml_a"].best_params["max_depth"] == tuned_params_a["max_depth"] - - # ensure tuning improved RMSE # not actually possible for ml_t as the targets are not available - assert tuned_score["ml_M"] < untuned_score["ml_M"] - assert tuned_score["ml_m"] < untuned_score["ml_m"] - assert tuned_score["ml_a"] < untuned_score["ml_a"] + dml_lplr.fit() + assert np.isfinite(dml_lplr.coef).all() + assert np.isfinite(dml_lplr.se).all() + + expected_a = np.empty_like(d) + for (_, test), inner_splits in zip(dml_lplr.smpls[0], dml_lplr.smpls_inner[0]): + predictions = [ + clone(tune_res[0]["ml_a"].best_estimator).fit(x[train], d[train]).predict(x[test]) for train, _ in inner_splits + ] + expected_a[test] = np.mean(predictions, axis=0) + np.testing.assert_allclose(dml_lplr.predictions["ml_a"][:, 0, 0], expected_a)