fix: carry over learners to the rebuilt sub-models in DoubleMLVector.sensitivity_benchmark

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)

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."""

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,
+    )