From 4cc84b417ae9ba33f53a9ae0f64359eb7de3f234 Mon Sep 17 00:00:00 2001 From: SvenKlaassen Date: Thu, 16 Oct 2025 10:14:29 +0200 Subject: [PATCH 1/5] update the set_ml_nuisance_params method to include the new parameter --- doubleml/double_ml.py | 24 +- doubleml/tests/test_set_ml_nuisance_params.py | 230 +++++++++++++----- .../test_set_ml_nuisance_params_models.py | 102 ++++++++ 3 files changed, 286 insertions(+), 70 deletions(-) create mode 100644 doubleml/tests/test_set_ml_nuisance_params_models.py diff --git a/doubleml/double_ml.py b/doubleml/double_ml.py index 8e7de835..5058951a 100644 --- a/doubleml/double_ml.py +++ b/doubleml/double_ml.py @@ -946,18 +946,34 @@ def set_ml_nuisance_params(self, learner, treat_var, params): ) if params is None: - all_params = [None] * self.n_rep + new_params = [None] * self.n_rep elif isinstance(params, dict): - all_params = [[params] * self.n_folds] * self.n_rep + new_params = [[params] * self.n_folds] * self.n_rep else: # ToDo: Add meaningful error message for asserts and corresponding uni tests assert len(params) == self.n_rep assert np.all(np.array([len(x) for x in params]) == self.n_folds) - all_params = params + new_params = params - self._params[learner][treat_var] = all_params + existing_params = self._params[learner].get(treat_var, [None] * self.n_rep) + if existing_params == [None] * self.n_rep: + updated_params = new_params + elif new_params == [None] * self.n_rep: + updated_params = existing_params + else: + updated_params = [] + for i_rep in range(self.n_rep): + rep_params = [] + for i_fold in range(self.n_folds): + existing_dict = existing_params[i_rep][i_fold] + new_dict = new_params[i_rep][i_fold] + updated_dict = existing_dict | new_dict + rep_params.append(updated_dict) + updated_params.append(rep_params) + + self._params[learner][treat_var] = updated_params return self @abstractmethod diff --git a/doubleml/tests/test_set_ml_nuisance_params.py b/doubleml/tests/test_set_ml_nuisance_params.py index 055bcbff..72cdceb2 100644 --- a/doubleml/tests/test_set_ml_nuisance_params.py +++ b/doubleml/tests/test_set_ml_nuisance_params.py @@ -2,101 +2,199 @@ import pytest from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor -from doubleml import DoubleMLCVAR, DoubleMLIIVM, DoubleMLIRM, DoubleMLLPQ, DoubleMLPLIV, DoubleMLPLR, DoubleMLPQ -from doubleml.irm.datasets import make_iivm_data, make_irm_data -from doubleml.plm.datasets import make_pliv_CHS2015, make_plr_CCDDHNR2018 +from doubleml import DoubleMLIRM, DoubleMLPLR +from doubleml.irm.datasets import make_irm_data +from doubleml.plm.datasets import make_plr_CCDDHNR2018 -# set default and test values -n_est_default = 100 -n_est_test = 5 -n_folds = 2 -test_values = [[{"n_estimators": 5}, {"n_estimators": 5}]] +# Test setup +n_folds = 3 +n_rep = 2 np.random.seed(3141) -dml_data_plr = make_plr_CCDDHNR2018(n_obs=100) -dml_data_pliv = make_pliv_CHS2015(n_obs=100, dim_z=1) dml_data_irm = make_irm_data(n_obs=1000) -dml_data_iivm = make_iivm_data(n_obs=2000) -reg_learner = RandomForestRegressor(max_depth=2) -class_learner = RandomForestClassifier(max_depth=2) +reg_learner = RandomForestRegressor(max_depth=2, n_estimators=100) +class_learner = RandomForestClassifier(max_depth=2, n_estimators=100) -# linear models -dml_plr = DoubleMLPLR(dml_data_plr, reg_learner, reg_learner, n_folds=n_folds) -dml_pliv = DoubleMLPLIV(dml_data_pliv, reg_learner, reg_learner, reg_learner, n_folds=n_folds) -dml_irm = DoubleMLIRM(dml_data_irm, reg_learner, class_learner, n_folds=n_folds) -dml_iivm = DoubleMLIIVM(dml_data_iivm, reg_learner, class_learner, class_learner, n_folds=n_folds) -dml_cvar = DoubleMLCVAR(dml_data_irm, ml_g=reg_learner, ml_m=class_learner, n_folds=n_folds) -dml_plr.set_ml_nuisance_params("ml_l", "d", {"n_estimators": n_est_test}) -dml_pliv.set_ml_nuisance_params("ml_l", "d", {"n_estimators": n_est_test}) -dml_irm.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test}) -dml_iivm.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test}) -dml_cvar.set_ml_nuisance_params("ml_g", "d", {"n_estimators": n_est_test}) +@pytest.fixture +def fresh_irm_model(): + """Create a fresh IRM model for each test.""" + return DoubleMLIRM(dml_data_irm, reg_learner, class_learner, n_folds=n_folds, n_rep=n_rep) -dml_plr.fit(store_models=True) -dml_pliv.fit(store_models=True) -dml_irm.fit(store_models=True) -dml_iivm.fit(store_models=True) -dml_cvar.fit(store_models=True) -# nonlinear models -dml_pq = DoubleMLPQ(dml_data_irm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds) -dml_lpq = DoubleMLLPQ(dml_data_iivm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds) +@pytest.mark.ci +def test_set_single_params(fresh_irm_model): + """Test combining behavior where new parameters are merged with existing ones.""" + dml_irm = fresh_irm_model + + # Set initial parameters + initial_params = {"n_estimators": 50, "max_depth": 3} + dml_irm.set_ml_nuisance_params("ml_g0", "d", initial_params) -dml_pq.set_ml_nuisance_params("ml_g", "d", {"n_estimators": n_est_test}) -dml_lpq.set_ml_nuisance_params("ml_m_z", "d", {"n_estimators": n_est_test}) + # Set additional parameters (should combine) + additional_params = {"min_samples_split": 5, "n_estimators": 25} # n_estimators should be updated + dml_irm.set_ml_nuisance_params("ml_g0", "d", additional_params) -dml_pq.fit(store_models=True) -dml_lpq.fit(store_models=True) + # With combining behavior, we should have all keys + expected_combined = {"n_estimators": 25, "max_depth": 3, "min_samples_split": 5} + assert dml_irm.params["ml_g0"]["d"][0][0] == expected_combined + assert dml_irm.params["ml_g0"]["d"][1][1] == expected_combined -def _assert_nuisance_params(dml_obj, learner_1, learner_2): - assert dml_obj.params[learner_1]["d"] == test_values - assert dml_obj.params[learner_2]["d"][0] is None +@pytest.mark.ci +def test_none_params_handling(fresh_irm_model): + """Test handling of None parameters.""" + dml_irm = fresh_irm_model - param_list_1 = [dml_obj.models[learner_1]["d"][0][fold].n_estimators for fold in range(n_folds)] - assert all(param == n_est_test for param in param_list_1) - param_list_2 = [dml_obj.models[learner_2]["d"][0][fold].n_estimators for fold in range(n_folds)] - assert all(param == n_est_default for param in param_list_2) + # Set initial parameters + initial_params = {"n_estimators": 50} + dml_irm.set_ml_nuisance_params("ml_g0", "d", initial_params) + # Setting None should not change existing parameters + dml_irm.set_ml_nuisance_params("ml_g0", "d", None) + assert dml_irm.params["ml_g0"]["d"][0][0] == initial_params -@pytest.mark.ci -def test_plr_params(): - _assert_nuisance_params(dml_plr, "ml_l", "ml_m") + # Test setting None on empty parameters + dml_irm.set_ml_nuisance_params("ml_g1", "d", None) + assert dml_irm.params["ml_g1"]["d"] == [None] * n_rep @pytest.mark.ci -def test_pliv_params(): - _assert_nuisance_params(dml_pliv, "ml_l", "ml_m") +def test_set_nested_list_params(fresh_irm_model): + """Test combining behavior with nested list parameters.""" + dml_irm = fresh_irm_model + + # Create initial nested parameters + initial_nested = [ + [ + {"n_estimators": 50, "max_depth": 2}, + {"n_estimators": 60, "max_depth": 3}, + {"n_estimators": 60, "max_depth": 3}, + ], # rep 0 + [ + {"n_estimators": 70, "max_depth": 4}, + {"n_estimators": 80, "max_depth": 5}, + {"n_estimators": 60, "max_depth": 3}, + ], # rep 1 + ] + dml_irm.set_ml_nuisance_params("ml_g0", "d", initial_nested) + + # Add additional parameters + additional_nested = [ + [ + {"min_samples_split": 2, "n_estimators": 25}, + {"min_samples_split": 3, "n_estimators": 35}, + {"min_samples_split": 3, "n_estimators": 35}, + ], # rep 0 + [ + {"min_samples_split": 4, "n_estimators": 45}, + {"min_samples_split": 5, "n_estimators": 55}, + {"min_samples_split": 3, "n_estimators": 35}, + ], # rep 1 + ] + dml_irm.set_ml_nuisance_params("ml_g0", "d", additional_nested) + + # Verify combining: existing keys preserved, overlapping keys updated, new keys added + expected_combined = [ + [ + {"n_estimators": 25, "max_depth": 2, "min_samples_split": 2}, + {"n_estimators": 35, "max_depth": 3, "min_samples_split": 3}, + {"n_estimators": 35, "max_depth": 3, "min_samples_split": 3}, + ], + [ + {"n_estimators": 45, "max_depth": 4, "min_samples_split": 4}, + {"n_estimators": 55, "max_depth": 5, "min_samples_split": 5}, + {"n_estimators": 35, "max_depth": 3, "min_samples_split": 3}, + ], + ] + + assert dml_irm.params["ml_g0"]["d"] == expected_combined @pytest.mark.ci -def test_irm_params(): - _assert_nuisance_params(dml_irm, "ml_g0", "ml_g1") +def test_multiple_learners_independence(fresh_irm_model): + """Test that parameters for different learners are independent.""" + dml_irm = fresh_irm_model + # Set parameters for different learners + params_g0 = {"n_estimators": 50} + params_g1 = {"n_estimators": 75} + params_m = {"n_estimators": 100} -@pytest.mark.ci -def test_iivm_params(): - _assert_nuisance_params(dml_iivm, "ml_g0", "ml_g1") + dml_irm.set_ml_nuisance_params("ml_g0", "d", params_g0) + dml_irm.set_ml_nuisance_params("ml_g1", "d", params_g1) + dml_irm.set_ml_nuisance_params("ml_m", "d", params_m) + # Verify independence + assert dml_irm.params["ml_g0"]["d"][0][0] == params_g0 + assert dml_irm.params["ml_g1"]["d"][0][0] == params_g1 + assert dml_irm.params["ml_m"]["d"][0][0] == params_m -@pytest.mark.ci -def test_cvar_params(): - _assert_nuisance_params(dml_cvar, "ml_g", "ml_m") + # Modify one learner, others should remain unchanged + new_params_g0 = {"max_depth": 3, "n_estimators": 25} + dml_irm.set_ml_nuisance_params("ml_g0", "d", new_params_g0) + + # With combining behavior + expected_g0 = {"n_estimators": 25, "max_depth": 3} + assert dml_irm.params["ml_g0"]["d"][0][0] == expected_g0 + assert dml_irm.params["ml_g1"]["d"][0][0] == params_g1 # unchanged + assert dml_irm.params["ml_m"]["d"][0][0] == params_m # unchanged @pytest.mark.ci -def test_pq_params(): - _assert_nuisance_params(dml_pq, "ml_g", "ml_m") +def test_multiple_treatment_variables_independence(): + """Test that parameters for different treatment variables are independent.""" + # Create PLR data with multiple treatment variables + np.random.seed(3141) + multi_treat_data = make_plr_CCDDHNR2018(n_obs=100) + + # Add a second treatment variable for testing + multi_treat_data.data["d2"] = np.random.normal(0, 1, 100) + multi_treat_data._d_cols = ["d", "d2"] + + dml_plr = DoubleMLPLR(multi_treat_data, reg_learner, reg_learner, n_folds=n_folds, n_rep=n_rep) + + # Set parameters for different treatment variables + params_d = {"n_estimators": 50} + params_d2 = {"n_estimators": 75} + + dml_plr.set_ml_nuisance_params("ml_l", "d", params_d) + dml_plr.set_ml_nuisance_params("ml_l", "d2", params_d2) + + # Verify independence + assert dml_plr.params["ml_l"]["d"][0][0] == params_d + assert dml_plr.params["ml_l"]["d2"][0][0] == params_d2 + + # Modify one treatment variable, other should remain unchanged + new_params_d = {"max_depth": 3, "n_estimators": 25} + dml_plr.set_ml_nuisance_params("ml_l", "d", new_params_d) + + # With combining behavior + expected_d = {"n_estimators": 25, "max_depth": 3} + assert dml_plr.params["ml_l"]["d"][0][0] == expected_d + assert dml_plr.params["ml_l"]["d2"][0][0] == params_d2 # unchanged @pytest.mark.ci -def test_lpq_params(): - _assert_nuisance_params(dml_lpq, "ml_m_z", "ml_m_d_z0") - param_list_2 = [dml_lpq.models["ml_m_d_z1"]["d"][0][fold].n_estimators for fold in range(n_folds)] - assert all(param == n_est_default for param in param_list_2) - param_list_2 = [dml_lpq.models["ml_g_du_z0"]["d"][0][fold].n_estimators for fold in range(n_folds)] - assert all(param == n_est_default for param in param_list_2) - param_list_2 = [dml_lpq.models["ml_g_du_z1"]["d"][0][fold].n_estimators for fold in range(n_folds)] - assert all(param == n_est_default for param in param_list_2) +def test_error_cases(fresh_irm_model): + """Test error handling for invalid inputs.""" + dml_irm = fresh_irm_model + + # Invalid learner + with pytest.raises(ValueError, match="Invalid nuisance learner"): + dml_irm.set_ml_nuisance_params("invalid_learner", "d", {"n_estimators": 50}) + + # Invalid treatment variable + with pytest.raises(ValueError, match="Invalid treatment variable"): + dml_irm.set_ml_nuisance_params("ml_g0", "invalid_treat", {"n_estimators": 50}) + + # Invalid nested list length (wrong n_rep) + invalid_nested = [[{"n_estimators": 50}, {"n_estimators": 60}]] # Only 1 rep, should be 2 + with pytest.raises(AssertionError): + dml_irm.set_ml_nuisance_params("ml_g0", "d", invalid_nested) + + # Invalid nested list length (wrong n_folds) + invalid_nested = [[{"n_estimators": 50}], [{"n_estimators": 60}]] # Only 1 fold, should be 2 # Only 1 fold, should be 2 + with pytest.raises(AssertionError): + dml_irm.set_ml_nuisance_params("ml_g0", "d", invalid_nested) diff --git a/doubleml/tests/test_set_ml_nuisance_params_models.py b/doubleml/tests/test_set_ml_nuisance_params_models.py new file mode 100644 index 00000000..055bcbff --- /dev/null +++ b/doubleml/tests/test_set_ml_nuisance_params_models.py @@ -0,0 +1,102 @@ +import numpy as np +import pytest +from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor + +from doubleml import DoubleMLCVAR, DoubleMLIIVM, DoubleMLIRM, DoubleMLLPQ, DoubleMLPLIV, DoubleMLPLR, DoubleMLPQ +from doubleml.irm.datasets import make_iivm_data, make_irm_data +from doubleml.plm.datasets import make_pliv_CHS2015, make_plr_CCDDHNR2018 + +# set default and test values +n_est_default = 100 +n_est_test = 5 +n_folds = 2 +test_values = [[{"n_estimators": 5}, {"n_estimators": 5}]] + +np.random.seed(3141) +dml_data_plr = make_plr_CCDDHNR2018(n_obs=100) +dml_data_pliv = make_pliv_CHS2015(n_obs=100, dim_z=1) +dml_data_irm = make_irm_data(n_obs=1000) +dml_data_iivm = make_iivm_data(n_obs=2000) + +reg_learner = RandomForestRegressor(max_depth=2) +class_learner = RandomForestClassifier(max_depth=2) + +# linear models +dml_plr = DoubleMLPLR(dml_data_plr, reg_learner, reg_learner, n_folds=n_folds) +dml_pliv = DoubleMLPLIV(dml_data_pliv, reg_learner, reg_learner, reg_learner, n_folds=n_folds) +dml_irm = DoubleMLIRM(dml_data_irm, reg_learner, class_learner, n_folds=n_folds) +dml_iivm = DoubleMLIIVM(dml_data_iivm, reg_learner, class_learner, class_learner, n_folds=n_folds) +dml_cvar = DoubleMLCVAR(dml_data_irm, ml_g=reg_learner, ml_m=class_learner, n_folds=n_folds) + +dml_plr.set_ml_nuisance_params("ml_l", "d", {"n_estimators": n_est_test}) +dml_pliv.set_ml_nuisance_params("ml_l", "d", {"n_estimators": n_est_test}) +dml_irm.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test}) +dml_iivm.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test}) +dml_cvar.set_ml_nuisance_params("ml_g", "d", {"n_estimators": n_est_test}) + +dml_plr.fit(store_models=True) +dml_pliv.fit(store_models=True) +dml_irm.fit(store_models=True) +dml_iivm.fit(store_models=True) +dml_cvar.fit(store_models=True) + +# nonlinear models +dml_pq = DoubleMLPQ(dml_data_irm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds) +dml_lpq = DoubleMLLPQ(dml_data_iivm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds) + +dml_pq.set_ml_nuisance_params("ml_g", "d", {"n_estimators": n_est_test}) +dml_lpq.set_ml_nuisance_params("ml_m_z", "d", {"n_estimators": n_est_test}) + +dml_pq.fit(store_models=True) +dml_lpq.fit(store_models=True) + + +def _assert_nuisance_params(dml_obj, learner_1, learner_2): + assert dml_obj.params[learner_1]["d"] == test_values + assert dml_obj.params[learner_2]["d"][0] is None + + param_list_1 = [dml_obj.models[learner_1]["d"][0][fold].n_estimators for fold in range(n_folds)] + assert all(param == n_est_test for param in param_list_1) + param_list_2 = [dml_obj.models[learner_2]["d"][0][fold].n_estimators for fold in range(n_folds)] + assert all(param == n_est_default for param in param_list_2) + + +@pytest.mark.ci +def test_plr_params(): + _assert_nuisance_params(dml_plr, "ml_l", "ml_m") + + +@pytest.mark.ci +def test_pliv_params(): + _assert_nuisance_params(dml_pliv, "ml_l", "ml_m") + + +@pytest.mark.ci +def test_irm_params(): + _assert_nuisance_params(dml_irm, "ml_g0", "ml_g1") + + +@pytest.mark.ci +def test_iivm_params(): + _assert_nuisance_params(dml_iivm, "ml_g0", "ml_g1") + + +@pytest.mark.ci +def test_cvar_params(): + _assert_nuisance_params(dml_cvar, "ml_g", "ml_m") + + +@pytest.mark.ci +def test_pq_params(): + _assert_nuisance_params(dml_pq, "ml_g", "ml_m") + + +@pytest.mark.ci +def test_lpq_params(): + _assert_nuisance_params(dml_lpq, "ml_m_z", "ml_m_d_z0") + param_list_2 = [dml_lpq.models["ml_m_d_z1"]["d"][0][fold].n_estimators for fold in range(n_folds)] + assert all(param == n_est_default for param in param_list_2) + param_list_2 = [dml_lpq.models["ml_g_du_z0"]["d"][0][fold].n_estimators for fold in range(n_folds)] + assert all(param == n_est_default for param in param_list_2) + param_list_2 = [dml_lpq.models["ml_g_du_z1"]["d"][0][fold].n_estimators for fold in range(n_folds)] + assert all(param == n_est_default for param in param_list_2) From ef71205011c89debbd30932480621e4b2a961b30 Mon Sep 17 00:00:00 2001 From: SvenKlaassen Date: Thu, 16 Oct 2025 10:24:47 +0200 Subject: [PATCH 2/5] add apo test --- doubleml/tests/test_set_ml_nuisance_params_models.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/doubleml/tests/test_set_ml_nuisance_params_models.py b/doubleml/tests/test_set_ml_nuisance_params_models.py index 055bcbff..c9717a77 100644 --- a/doubleml/tests/test_set_ml_nuisance_params_models.py +++ b/doubleml/tests/test_set_ml_nuisance_params_models.py @@ -2,7 +2,7 @@ import pytest from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor -from doubleml import DoubleMLCVAR, DoubleMLIIVM, DoubleMLIRM, DoubleMLLPQ, DoubleMLPLIV, DoubleMLPLR, DoubleMLPQ +from doubleml import DoubleMLAPO, DoubleMLCVAR, DoubleMLIIVM, DoubleMLIRM, DoubleMLLPQ, DoubleMLPLIV, DoubleMLPLR, DoubleMLPQ from doubleml.irm.datasets import make_iivm_data, make_irm_data from doubleml.plm.datasets import make_pliv_CHS2015, make_plr_CCDDHNR2018 @@ -27,18 +27,21 @@ dml_irm = DoubleMLIRM(dml_data_irm, reg_learner, class_learner, n_folds=n_folds) dml_iivm = DoubleMLIIVM(dml_data_iivm, reg_learner, class_learner, class_learner, n_folds=n_folds) dml_cvar = DoubleMLCVAR(dml_data_irm, ml_g=reg_learner, ml_m=class_learner, n_folds=n_folds) +dml_apo = DoubleMLAPO(dml_data_irm, ml_g=reg_learner, ml_m=class_learner, n_folds=n_folds, treatment_level=1) dml_plr.set_ml_nuisance_params("ml_l", "d", {"n_estimators": n_est_test}) dml_pliv.set_ml_nuisance_params("ml_l", "d", {"n_estimators": n_est_test}) dml_irm.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test}) dml_iivm.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test}) dml_cvar.set_ml_nuisance_params("ml_g", "d", {"n_estimators": n_est_test}) +dml_apo.set_ml_nuisance_params("ml_g_d_lvl1", "d", {"n_estimators": n_est_test}) dml_plr.fit(store_models=True) dml_pliv.fit(store_models=True) dml_irm.fit(store_models=True) dml_iivm.fit(store_models=True) dml_cvar.fit(store_models=True) +dml_apo.fit(store_models=True) # nonlinear models dml_pq = DoubleMLPQ(dml_data_irm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds) @@ -76,6 +79,11 @@ def test_irm_params(): _assert_nuisance_params(dml_irm, "ml_g0", "ml_g1") +@pytest.mark.ci +def test_apo_params(): + _assert_nuisance_params(dml_apo, "ml_g_d_lvl1", "ml_m") + + @pytest.mark.ci def test_iivm_params(): _assert_nuisance_params(dml_iivm, "ml_g0", "ml_g1") From 92b1f7079ae91134d37e2b19f5845ccf57f304f0 Mon Sep 17 00:00:00 2001 From: SvenKlaassen Date: Thu, 16 Oct 2025 10:52:06 +0200 Subject: [PATCH 3/5] add did binary models to set nuisance params test --- .../test_set_ml_nuisance_params_models.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/doubleml/tests/test_set_ml_nuisance_params_models.py b/doubleml/tests/test_set_ml_nuisance_params_models.py index c9717a77..fd4b1639 100644 --- a/doubleml/tests/test_set_ml_nuisance_params_models.py +++ b/doubleml/tests/test_set_ml_nuisance_params_models.py @@ -3,6 +3,9 @@ from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor from doubleml import DoubleMLAPO, DoubleMLCVAR, DoubleMLIIVM, DoubleMLIRM, DoubleMLLPQ, DoubleMLPLIV, DoubleMLPLR, DoubleMLPQ +from doubleml.data import DoubleMLPanelData +from doubleml.did import DoubleMLDIDBinary, DoubleMLDIDCSBinary +from doubleml.did.datasets import make_did_CS2021 from doubleml.irm.datasets import make_iivm_data, make_irm_data from doubleml.plm.datasets import make_pliv_CHS2015, make_plr_CCDDHNR2018 @@ -18,6 +21,10 @@ dml_data_irm = make_irm_data(n_obs=1000) dml_data_iivm = make_iivm_data(n_obs=2000) +# Create DID data +df_did = make_did_CS2021(n_obs=500, dgp_type=1, n_pre_treat_periods=2, n_periods=4, time_type="float") +dml_data_did = DoubleMLPanelData(df_did, y_col="y", d_cols="d", id_col="id", t_col="t", x_cols=["Z1", "Z2", "Z3", "Z4"]) + reg_learner = RandomForestRegressor(max_depth=2) class_learner = RandomForestClassifier(max_depth=2) @@ -43,6 +50,35 @@ dml_cvar.fit(store_models=True) dml_apo.fit(store_models=True) +# DID models +dml_did_binary = DoubleMLDIDBinary( + obj_dml_data=dml_data_did, + ml_g=reg_learner, + ml_m=class_learner, + g_value=2, + t_value_pre=0, + t_value_eval=1, + score="observational", + n_folds=n_folds, +) + +dml_did_cs_binary = DoubleMLDIDCSBinary( + obj_dml_data=dml_data_did, + ml_g=reg_learner, + ml_m=class_learner, + g_value=2, + t_value_pre=0, + t_value_eval=1, + score="observational", + n_folds=n_folds, +) + +dml_did_binary.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test}) +dml_did_cs_binary.set_ml_nuisance_params("ml_g_d0_t0", "d", {"n_estimators": n_est_test}) + +dml_did_binary.fit(store_models=True) +dml_did_cs_binary.fit(store_models=True) + # nonlinear models dml_pq = DoubleMLPQ(dml_data_irm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds) dml_lpq = DoubleMLLPQ(dml_data_iivm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds) @@ -108,3 +144,13 @@ def test_lpq_params(): assert all(param == n_est_default for param in param_list_2) param_list_2 = [dml_lpq.models["ml_g_du_z1"]["d"][0][fold].n_estimators for fold in range(n_folds)] assert all(param == n_est_default for param in param_list_2) + + +@pytest.mark.ci +def test_did_binary_params(): + _assert_nuisance_params(dml_did_binary, "ml_g0", "ml_g1") + + +@pytest.mark.ci +def test_did_cs_binary_params(): + _assert_nuisance_params(dml_did_cs_binary, "ml_g_d0_t0", "ml_g_d1_t0") From cdbda92afeb2e5c70053f606a30dc4ccfbd7d185 Mon Sep 17 00:00:00 2001 From: SvenKlaassen Date: Thu, 16 Oct 2025 11:42:45 +0200 Subject: [PATCH 4/5] add testt for tuning with set parameters --- doubleml/tests/test_set_ml_nuisance_params.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/doubleml/tests/test_set_ml_nuisance_params.py b/doubleml/tests/test_set_ml_nuisance_params.py index 72cdceb2..6c8c19a3 100644 --- a/doubleml/tests/test_set_ml_nuisance_params.py +++ b/doubleml/tests/test_set_ml_nuisance_params.py @@ -198,3 +198,56 @@ def test_error_cases(fresh_irm_model): invalid_nested = [[{"n_estimators": 50}], [{"n_estimators": 60}]] # Only 1 fold, should be 2 # Only 1 fold, should be 2 with pytest.raises(AssertionError): dml_irm.set_ml_nuisance_params("ml_g0", "d", invalid_nested) + + +@pytest.mark.ci +def test_set_params_then_tune_combination(fresh_irm_model): + """Test that manually set parameters are preserved and combined with tuned parameters.""" + dml_irm = fresh_irm_model + + # Set initial parameters that should be preserved after tuning + initial_params = {"max_depth": 3, "min_samples_split": 5} + dml_irm.set_ml_nuisance_params("ml_g0", "d", initial_params) + dml_irm.set_ml_nuisance_params("ml_g1", "d", initial_params) + dml_irm.set_ml_nuisance_params("ml_m", "d", {"max_depth": 2}) + + # Define tuning grid - only tune n_estimators, min_samples_split, not all manually set parameters + par_grid = {"ml_g": {"n_estimators": [10, 20], "min_samples_split": [2, 10]}, "ml_m": {"n_estimators": [15, 25]}} + tune_res = dml_irm.tune(par_grid, return_tune_res=True) + + # Verify consistency across folds and repetitions + for rep in range(n_rep): + for fold in range(n_folds): + # All should have the same combination of manually set + tuned parameters + fold_g0_params = dml_irm.params["ml_g0"]["d"][rep][fold] + fold_g1_params = dml_irm.params["ml_g1"]["d"][rep][fold] + fold_m_params = dml_irm.params["ml_m"]["d"][rep][fold] + + # Manually set parameters that are not tuned should be preserved + assert fold_g0_params["max_depth"] == 3 + assert fold_g1_params["max_depth"] == 3 + assert fold_m_params["max_depth"] == 2 + + # Tuned parameters should overwrite manually set ones + assert fold_g0_params["n_estimators"] in [10, 20] + assert fold_g1_params["n_estimators"] in [10, 20] + assert fold_m_params["n_estimators"] in [15, 25] + + # min_samples_split should be overwritten by tuning for ml_g learners + assert fold_g0_params["min_samples_split"] in [2, 10] + assert fold_g1_params["min_samples_split"] in [2, 10] + + # Check that manually set max_depth is preserved in best estimators + for fold in range(n_folds): + # Check if tune_res contains GridSearchCV objects + if hasattr(tune_res[0]["tune_res"]["g0_tune"][fold], "best_estimator_"): + best_estimator_g0 = tune_res[0]["tune_res"]["g0_tune"][fold].best_estimator_ + assert best_estimator_g0.max_depth == 3 + + if hasattr(tune_res[0]["tune_res"]["g1_tune"][fold], "best_estimator_"): + best_estimator_g1 = tune_res[0]["tune_res"]["g1_tune"][fold].best_estimator_ + assert best_estimator_g1.max_depth == 3 + + if hasattr(tune_res[0]["tune_res"]["m_tune"][fold], "best_estimator_"): + best_estimator_m = tune_res[0]["tune_res"]["m_tune"][fold].best_estimator_ + assert best_estimator_m.max_depth == 2 From f8a7369f68fa76781a65222fc0e3634c204b87ce Mon Sep 17 00:00:00 2001 From: SvenKlaassen Date: Thu, 16 Oct 2025 13:00:01 +0200 Subject: [PATCH 5/5] update test --- doubleml/tests/test_set_ml_nuisance_params.py | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/doubleml/tests/test_set_ml_nuisance_params.py b/doubleml/tests/test_set_ml_nuisance_params.py index 6c8c19a3..f338b3aa 100644 --- a/doubleml/tests/test_set_ml_nuisance_params.py +++ b/doubleml/tests/test_set_ml_nuisance_params.py @@ -213,7 +213,7 @@ def test_set_params_then_tune_combination(fresh_irm_model): # Define tuning grid - only tune n_estimators, min_samples_split, not all manually set parameters par_grid = {"ml_g": {"n_estimators": [10, 20], "min_samples_split": [2, 10]}, "ml_m": {"n_estimators": [15, 25]}} - tune_res = dml_irm.tune(par_grid, return_tune_res=True) + dml_irm.tune(par_grid, return_tune_res=False) # Verify consistency across folds and repetitions for rep in range(n_rep): @@ -236,18 +236,3 @@ def test_set_params_then_tune_combination(fresh_irm_model): # min_samples_split should be overwritten by tuning for ml_g learners assert fold_g0_params["min_samples_split"] in [2, 10] assert fold_g1_params["min_samples_split"] in [2, 10] - - # Check that manually set max_depth is preserved in best estimators - for fold in range(n_folds): - # Check if tune_res contains GridSearchCV objects - if hasattr(tune_res[0]["tune_res"]["g0_tune"][fold], "best_estimator_"): - best_estimator_g0 = tune_res[0]["tune_res"]["g0_tune"][fold].best_estimator_ - assert best_estimator_g0.max_depth == 3 - - if hasattr(tune_res[0]["tune_res"]["g1_tune"][fold], "best_estimator_"): - best_estimator_g1 = tune_res[0]["tune_res"]["g1_tune"][fold].best_estimator_ - assert best_estimator_g1.max_depth == 3 - - if hasattr(tune_res[0]["tune_res"]["m_tune"][fold], "best_estimator_"): - best_estimator_m = tune_res[0]["tune_res"]["m_tune"][fold].best_estimator_ - assert best_estimator_m.max_depth == 2