diff --git a/doubleml/double_ml.py b/doubleml/double_ml.py index 22a066cc..f47273c5 100644 --- a/doubleml/double_ml.py +++ b/doubleml/double_ml.py @@ -895,7 +895,7 @@ def tune( raise TypeError(f"return_tune_res must be True or False. Got {str(return_tune_res)}.") if tune_on_folds: - tuning_res = [[None] * self.n_rep] * self._dml_data.n_treat + tuning_res = [[None] * self.n_rep for _ in range(self._dml_data.n_treat)] else: tuning_res = [None] * self._dml_data.n_treat @@ -921,7 +921,7 @@ def tune( n_iter_randomized_search, ) - tuning_res[i_rep][i_d] = res + tuning_res[i_d][i_rep] = res nuisance_params.append(res["params"]) if set_as_params: diff --git a/doubleml/tests/test_tune_on_folds.py b/doubleml/tests/test_tune_on_folds.py new file mode 100644 index 00000000..a805f845 --- /dev/null +++ b/doubleml/tests/test_tune_on_folds.py @@ -0,0 +1,30 @@ +import numpy as np +import pytest +from sklearn.linear_model import Lasso + +import doubleml as dml +from doubleml.plm.datasets import make_plr_CCDDHNR2018 + + +@pytest.mark.ci +@pytest.mark.parametrize("n_rep, n_treat", [(1, 1), (2, 1), (1, 2), (3, 2)]) +def test_tune_on_folds_result_per_treatment_and_repetition(n_rep, n_treat): + np.random.seed(3141) + df = make_plr_CCDDHNR2018(n_obs=200, dim_x=5, return_type="DataFrame") + d_cols = ["d", "X1"][:n_treat] + x_cols = [c for c in df.columns if c.startswith("X") and c not in d_cols] + dml_data = dml.DoubleMLData(df, "y", d_cols, x_cols) + dml_plr = dml.DoubleMLPLR(dml_data, Lasso(), Lasso(), n_folds=2, n_rep=n_rep) + + grid = {"ml_l": {"alpha": [0.05, 0.1]}, "ml_m": {"alpha": [0.05, 0.1]}} + tune_res = dml_plr.tune(grid, tune_on_folds=True, n_folds_tune=2, return_tune_res=True) + + assert len(tune_res) == n_treat + for i_d in range(n_treat): + assert len(tune_res[i_d]) == n_rep + assert len({id(res) for res in tune_res[i_d]}) == n_rep + for i_rep in range(n_rep): + # the results of each treatment and repetition are the parameters that were set + params = dml_plr.params["ml_l"][d_cols[i_d]][i_rep] + assert tune_res[i_d][i_rep]["params"]["ml_l"] == params + dml_plr.fit()