Pass feature set to _evaluate_model_kind

Quantco · Jul 19, 2024 · d81ea7a · d81ea7a
1 parent f66a7f0
commit d81ea7a
Show file tree

Hide file tree

Showing 7 changed files with 17 additions and 3 deletions.
diff --git a/metalearners/_typing.py b/metalearners/_typing.py
@@ -17,7 +17,7 @@
 OosMethod = Literal["overall", "median", "mean"]
 
 Params = Mapping[str, int | float | str]
-Features = Collection[str] | Collection[int]
+Features = Collection[str] | Collection[int] | None
 
 # ruff is not happy about the usage of Union.
 Vector = Union[pd.Series, np.ndarray]  # noqa

diff --git a/metalearners/drlearner.py b/metalearners/drlearner.py
@@ -267,6 +267,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
         )
 
         propensity_evaluation = _evaluate_model_kind(
@@ -278,6 +279,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[PROPENSITY_MODEL],
         )
 
         pseudo_outcome: list[np.ndarray] = []
@@ -301,6 +303,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=True,
+            feature_set=self.feature_set[TREATMENT_MODEL],
         )
 
         return variant_outcome_evaluation | propensity_evaluation | treatment_evaluation

diff --git a/metalearners/metalearner.py b/metalearners/metalearner.py
@@ -145,6 +145,7 @@ def _evaluate_model_kind(
     model_kind: str,
     is_oos: bool,
     is_treatment_model: bool,
+    feature_set: Features,
     oos_method: OosMethod = OVERALL,
     sample_weights: Sequence[Vector] | None = None,
 ) -> dict[str, float]:
@@ -168,14 +169,15 @@ def _evaluate_model_kind(
                 else:
                     index_str = f"{i}_"
             name = f"{prefix}{index_str}{scorer_name}"
+            X_filtered = _filter_x_columns(Xs[i], feature_set)
             with _PredictContext(cfe, is_oos, oos_method) as modified_cfe:
                 if sample_weights:
                     evaluation_metrics[name] = scorer_callable(
-                        modified_cfe, Xs[i], ys[i], sample_weight=sample_weights[i]
+                        modified_cfe, X_filtered, ys[i], sample_weight=sample_weights[i]
                     )
                 else:
                     evaluation_metrics[name] = scorer_callable(
-                        modified_cfe, Xs[i], ys[i]
+                        modified_cfe, X_filtered, ys[i]
                     )
     return evaluation_metrics
 

diff --git a/metalearners/rlearner.py b/metalearners/rlearner.py
@@ -352,6 +352,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[PROPENSITY_MODEL],
         )
 
         outcome_evaluation = _evaluate_model_kind(
@@ -363,6 +364,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[OUTCOME_MODEL],
         )
 
         # TODO: improve this? generalize it to other metalearners?
@@ -414,6 +416,7 @@ def evaluate(
             oos_method=oos_method,
             is_treatment_model=True,
             sample_weights=sample_weights,
+            feature_set=self.feature_set[TREATMENT_MODEL],
         )
 
         rloss_evaluation = {}

diff --git a/metalearners/slearner.py b/metalearners/slearner.py
@@ -212,6 +212,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[_BASE_MODEL],
         )
 
     def predict_conditional_average_outcomes(

diff --git a/metalearners/tlearner.py b/metalearners/tlearner.py
@@ -126,4 +126,5 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
         )
diff --git a/metalearners/xlearner.py b/metalearners/xlearner.py
@@ -300,6 +300,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
         )
 
         propensity_evaluation = _evaluate_model_kind(
@@ -311,6 +312,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=False,
+            feature_set=self.feature_set[PROPENSITY_MODEL],
         )
 
         imputed_te_control: list[np.ndarray] = []
@@ -331,6 +333,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=True,
+            feature_set=self.feature_set[TREATMENT_EFFECT_MODEL],
         )
 
         te_control_evaluation = _evaluate_model_kind(
@@ -342,6 +345,7 @@ def evaluate(
             is_oos=is_oos,
             oos_method=oos_method,
             is_treatment_model=True,
+            feature_set=self.feature_set[CONTROL_EFFECT_MODEL],
         )
 
         return (