From 5b4817cc41872c4dc30770ae08bc827d9a351f9a Mon Sep 17 00:00:00 2001 From: Rajesh Vakkalagadda Date: Mon, 2 Nov 2026 11:22:45 -0800 Subject: [PATCH 1/3] fix: carry over learners to the rebuilt sub-models in DoubleMLVector.sensitivity_benchmark DoubleMLVector.sensitivity_benchmark rebuilt the sub-models of the short-form model with _initialize_models() so that they use the reduced covariate set, but the new sub-models had no learners: in the vector API, the learners are only stored on the sub-models. Refitting the short-form model therefore always raised "Learner 'ml_l' is required but not set". Give each rebuilt sub-model the learners of the sub-model it replaces. Check that PLRVector.sensitivity_benchmark matches the legacy DoubleMLPLR.sensitivity_benchmark. --- doubleml/double_ml_vector.py | 6 ++++++ doubleml/plm/tests/test_plr_vector_vs_plr.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/doubleml/double_ml_vector.py b/doubleml/double_ml_vector.py index 3990dff9..dac764e9 100644 --- a/doubleml/double_ml_vector.py +++ b/doubleml/double_ml_vector.py @@ -742,9 +742,15 @@ def sensitivity_benchmark(self, benchmarking_set: list[str], fit_args: dict | No x_list_short = [x for x in x_list_long if x not in benchmarking_set] dml_short = copy.deepcopy(self) dml_short._dml_data.x_cols = x_list_short + long_models = dml_short._modellist + if long_models is None: + raise ValueError("Sub-models are not initialized.") # Sub-models each hold their own DoubleMLData — rebuild them from the updated _dml_data # so that the short-form model actually uses the reduced feature set. dml_short._modellist = dml_short._initialize_models() + # Learners are only stored on the sub-models, so carry them over to the rebuilt ones. + for model_short, model_long in zip(dml_short._modellist, long_models): + model_short._learners = model_long._learners dml_short._framework = None if fit_args is not None: diff --git a/doubleml/plm/tests/test_plr_vector_vs_plr.py b/doubleml/plm/tests/test_plr_vector_vs_plr.py index 58e28d4a..9c9c9982 100644 --- a/doubleml/plm/tests/test_plr_vector_vs_plr.py +++ b/doubleml/plm/tests/test_plr_vector_vs_plr.py @@ -1,6 +1,7 @@ """Compare PLRVector against the legacy DoubleMLPLR implementation for multi-treatment data.""" import numpy as np +import pandas as pd import pytest from sklearn.base import clone from sklearn.linear_model import Lasso, LinearRegression @@ -120,3 +121,16 @@ def test_sensitivity_max_bias_equal(comparison_fixture): old.framework.sensitivity_elements["max_bias"], rtol=1e-9, ) + + +@pytest.mark.ci +def test_sensitivity_benchmark_equal(comparison_fixture): + """PLRVector.sensitivity_benchmark matches legacy DoubleMLPLR.sensitivity_benchmark.""" + old = comparison_fixture["old"] + new = comparison_fixture["new"] + benchmarking_set = ["X1"] + pd.testing.assert_frame_equal( + new.sensitivity_benchmark(benchmarking_set=benchmarking_set), + old.sensitivity_benchmark(benchmarking_set=benchmarking_set), + rtol=1e-9, + ) From 00dba3ceb7b5ab470c786bef5629f4cb6b749439 Mon Sep 17 00:00:00 2001 From: SvenKlaassen Date: Tue, 3 Nov 2026 07:29:46 +0100 Subject: [PATCH 2/3] add _rebuild_models method to retain learners when recreating sub-models --- doubleml/double_ml_vector.py | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/doubleml/double_ml_vector.py b/doubleml/double_ml_vector.py index dac764e9..810cb863 100644 --- a/doubleml/double_ml_vector.py +++ b/doubleml/double_ml_vector.py @@ -351,6 +351,29 @@ def _reset_fit_state(self) -> None: for model in self._modellist: model._reset_fit_state() + def _rebuild_models(self) -> None: + """ + Re-create the sub-models from the current ``_dml_data``, keeping their learners. + + Each sub-model holds its own data view, so changes to ``_dml_data`` (e.g. ``x_cols``) + only take effect after the sub-models are rebuilt. Learners are only stored on the + sub-models and are carried over by position. Fit state is cleared. + + Raises + ------ + ValueError + If the sub-models are not initialized or the rebuilt sub-models do not match + the existing ones in number. + + """ + if self._modellist is None: + raise ValueError("Sub-models are not initialized. Call _initialize_models() in the subclass __init__.") + new_models = self._initialize_models() + for new_model, old_model in zip(new_models, self._modellist, strict=True): + new_model._learners = old_model._learners + self._modellist = new_models + self._framework = None + def _propagate_splits_to_models(self) -> None: """Push the vector's sample splits into each sub-model.""" if self._modellist is None: @@ -742,16 +765,7 @@ def sensitivity_benchmark(self, benchmarking_set: list[str], fit_args: dict | No x_list_short = [x for x in x_list_long if x not in benchmarking_set] dml_short = copy.deepcopy(self) dml_short._dml_data.x_cols = x_list_short - long_models = dml_short._modellist - if long_models is None: - raise ValueError("Sub-models are not initialized.") - # Sub-models each hold their own DoubleMLData — rebuild them from the updated _dml_data - # so that the short-form model actually uses the reduced feature set. - dml_short._modellist = dml_short._initialize_models() - # Learners are only stored on the sub-models, so carry them over to the rebuilt ones. - for model_short, model_long in zip(dml_short._modellist, long_models): - model_short._learners = model_long._learners - dml_short._framework = None + dml_short._rebuild_models() if fit_args is not None: dml_short.fit(**fit_args) From 5d4da666a68c9cfce5430d0dd059c6ab60662a7f Mon Sep 17 00:00:00 2001 From: SvenKlaassen Date: Tue, 3 Nov 2026 07:56:16 +0100 Subject: [PATCH 3/3] extend PLRVector sensitivity_benchmark coverage --- .../plm/tests/test_plr_vector_exceptions.py | 20 +++++++++++++ doubleml/plm/tests/test_plr_vector_vs_plr.py | 29 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/doubleml/plm/tests/test_plr_vector_exceptions.py b/doubleml/plm/tests/test_plr_vector_exceptions.py index cc4fdff3..b1d71b30 100644 --- a/doubleml/plm/tests/test_plr_vector_exceptions.py +++ b/doubleml/plm/tests/test_plr_vector_exceptions.py @@ -151,6 +151,26 @@ def test_warning_ml_g_partialling_out(): dml_obj.set_learners(ml_l=Lasso(alpha=0.1), ml_m=Lasso(alpha=0.1), ml_g=Lasso(alpha=0.1)) +@pytest.mark.ci +@pytest.mark.parametrize( + "kwargs, error, msg", + [ + ({"benchmarking_set": "X1"}, TypeError, r"benchmarking_set must be a list\."), + ({"benchmarking_set": []}, ValueError, r"benchmarking_set must not be empty\."), + ({"benchmarking_set": ["d2"]}, ValueError, r"benchmarking_set must be a subset of features"), + ({"benchmarking_set": ["X1"], "fit_args": 1}, TypeError, r"fit_args must be a dict\."), + ], +) +def test_exception_sensitivity_benchmark(kwargs, error, msg): + """sensitivity_benchmark rejects calls before fit() and invalid inputs.""" + dml_obj = PLRVector(_make_bivariate_data(), ml_l=Lasso(alpha=0.1), ml_m=Lasso(alpha=0.1)) + with pytest.raises(ValueError, match=r"Apply fit\(\) before sensitivity_benchmark\(\)\."): + dml_obj.sensitivity_benchmark(benchmarking_set=["X1"]) + dml_obj.fit(n_folds=3) + with pytest.raises(error, match=msg): + dml_obj.sensitivity_benchmark(**kwargs) + + @pytest.mark.ci def test_cate_not_implemented(): """cate() raises NotImplementedError on multi-treatment PLR.""" diff --git a/doubleml/plm/tests/test_plr_vector_vs_plr.py b/doubleml/plm/tests/test_plr_vector_vs_plr.py index 9c9c9982..a4ace094 100644 --- a/doubleml/plm/tests/test_plr_vector_vs_plr.py +++ b/doubleml/plm/tests/test_plr_vector_vs_plr.py @@ -129,8 +129,37 @@ def test_sensitivity_benchmark_equal(comparison_fixture): old = comparison_fixture["old"] new = comparison_fixture["new"] benchmarking_set = ["X1"] + coef, x_cols = new.coef.copy(), list(new._dml_data.x_cols) pd.testing.assert_frame_equal( new.sensitivity_benchmark(benchmarking_set=benchmarking_set), old.sensitivity_benchmark(benchmarking_set=benchmarking_set), rtol=1e-9, ) + # The short-form refit must not modify the long-form model + np.testing.assert_array_equal(new.coef, coef) + assert list(new._dml_data.x_cols) == x_cols + + +@pytest.mark.ci +def test_sensitivity_benchmark_per_treatment_params(generate_data_bivariate): + """Treatment-specific learner params survive the short-form rebuild (matches legacy).""" + data = generate_data_bivariate + x_cols = data.columns[data.columns.str.startswith("X")].tolist() + obj_dml_data = dml.DoubleMLData(data, y_col="y", d_cols=["d1", "d2"], x_cols=x_cols) + + np.random.seed(3141) + old = dml.DoubleMLPLR(obj_dml_data, Lasso(alpha=0.1), Lasso(alpha=0.1), n_folds=3) + old.set_ml_nuisance_params("ml_l", "d2", {"alpha": 0.5}) + old.fit() + + np.random.seed(3141) + new = PLRVector(obj_dml_data, ml_l=Lasso(alpha=0.1), ml_m=Lasso(alpha=0.1)) + new.draw_sample_splitting(n_folds=3) + new.modellist[1].set_params("ml_l", alpha=0.5) + new.fit() + + pd.testing.assert_frame_equal( + new.sensitivity_benchmark(benchmarking_set=["X1"]), + old.sensitivity_benchmark(benchmarking_set=["X1"]), + rtol=1e-9, + )