diff --git a/doubleml/double_ml_vector.py b/doubleml/double_ml_vector.py index 3990dff9..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,10 +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 - # 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() - dml_short._framework = None + dml_short._rebuild_models() if fit_args is not None: dml_short.fit(**fit_args) 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 58e28d4a..a4ace094 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,45 @@ 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"] + 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, + )