Raise error if fitted propensity is passed through nuisance models

metalearners/metalearner.py

@@ -451,10 +451,10 @@ def __init__(
         if fitted_nuisance_models is not None:
             if not set(fitted_nuisance_models.keys()) <= set(
                 nuisance_model_specifications.keys()
-            ):
+            ) - {PROPENSITY_MODEL}:
                 raise ValueError(
-                    "The keys present in fitted_nuisance_models should be a subset of"
-                    f"{set(nuisance_model_specifications.keys())}"
+                    "The keys present in fitted_nuisance_models should be a subset of "
+                    f"{set(nuisance_model_specifications.keys()) - {PROPENSITY_MODEL}}"
                 )
             self._nuisance_models |= deepcopy(fitted_nuisance_models)
             not_fitted_nuisance_models -= set(fitted_nuisance_models.keys())

tests/test_metalearner.py

@@ -470,6 +470,7 @@ def test_feature_set(feature_set, expected_n_features, use_pandas, rng):
 
 
 def test_model_reusage_init():
+    # TODO: Split up into several tests.
     prefitted_models = [CrossFitEstimator(10, LGBMRegressor)]
     ml = _TestMetaLearner(
         nuisance_model_factory=LinearRegression,
@@ -497,6 +498,16 @@ def test_model_reusage_init():
             fitted_nuisance_models={"nuisance3": prefitted_models},
         )
 
+    with pytest.raises(ValueError, match="The keys present"):
+        RLearner(
+            propensity_model_factory=LGBMClassifier,
+            nuisance_model_factory=LGBMRegressor,
+            is_classification=False,
+            n_variants=2,
+            treatment_model_factory=LGBMRegressor,
+            fitted_nuisance_models={PROPENSITY_MODEL: prefitted_models},
+        )
+
 
 @pytest.mark.parametrize(
     "fit_params, nuisance_model_names, treatment_model_names, expected",