Merge branch 'main' into survival_example

CHANGELOG.rst

@@ -7,6 +7,17 @@
 Changelog
 =========
 
+0.8.0 (2024-07-xx)
+------------------
+
+* Implement :meth:`metalearners.cross_fit_estimator.CrossFitEstimator.score`.
+
+**Bug fixes**
+
+* Fixed a bug in :meth:`metalearners.metalearner.MetaLearner.evaluate` where it failed
+  in the case of ``feature_set`` being different from ``None``.
+
+
 0.7.0 (2024-07-12)
 ------------------
 

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

metalearners/cross_fit_estimator.py

@@ -6,7 +6,8 @@
 from functools import partial
 
 import numpy as np
-from sklearn.base import is_classifier
+from sklearn.base import is_classifier, is_regressor
+from sklearn.metrics import accuracy_score, r2_score
 from sklearn.model_selection import (
     KFold,
     StratifiedKFold,
@@ -337,8 +338,28 @@ def predict_proba(
             oos_method=oos_method,
         )
 
-    def score(self, X, y, sample_weight=None, **kwargs):
-        raise NotImplementedError()
+    def score(
+        self,
+        X: Matrix,
+        y: Vector,
+        is_oos: bool,
+        oos_method: OosMethod | None = None,
+        sample_weight: Vector | None = None,
+    ) -> float:
+        """Return the coefficient of determination of the prediction if the estimator is
+        a regressor or the mean accuracy if it is a classifier."""
+        if is_classifier(self):
+            return accuracy_score(
+                y, self.predict(X, is_oos, oos_method), sample_weight=sample_weight
+            )
+        elif is_regressor(self):
+            return r2_score(
+                y, self.predict(X, is_oos, oos_method), sample_weight=sample_weight
+            )
+        else:
+            raise NotImplementedError(
+                "score is not implemented for this type of estimator."
+            )
 
     def set_params(self, **params):
         raise NotImplementedError()

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

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
 

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 = {}

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(

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

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 (

tests/test_cross_fit_estimator.py

@@ -6,6 +6,7 @@
 import numpy as np
 import pytest
 from lightgbm import LGBMClassifier, LGBMRegressor
+from sklearn.base import is_classifier, is_regressor
 from sklearn.linear_model import LinearRegression, LogisticRegression
 from sklearn.metrics import accuracy_score, log_loss
 from sklearn.model_selection import KFold
@@ -262,3 +263,20 @@ def test_validate_data_match(n_observations, test_indices, success):
             ValueError, match="rely on different numbers of observations"
         ):
             _validate_data_match_prior_split(n_observations, test_indices)
+
+
+@pytest.mark.parametrize(
+    "estimator",
+    [LGBMClassifier, LGBMRegressor],
+)
+def test_score_smoke(estimator, rng):
+    n_samples = 1000
+    X = rng.standard_normal((n_samples, 3))
+    if is_classifier(estimator):
+        y = rng.integers(0, 4, n_samples)
+    elif is_regressor(estimator):
+        y = rng.standard_normal(n_samples)
+
+    cfe = CrossFitEstimator(5, estimator, {"n_estimators": 3})
+    cfe.fit(X, y)
+    cfe.score(X, y, False)

tests/test_learner.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import numpy as np
+import pandas as pd
 import pytest
 from lightgbm import LGBMClassifier, LGBMRegressor
 from sklearn.linear_model import LinearRegression, LogisticRegression
@@ -12,13 +13,15 @@
 from metalearners.drlearner import DRLearner
 from metalearners.metalearner import (
     PROPENSITY_MODEL,
+    TREATMENT_MODEL,
     VARIANT_OUTCOME_MODEL,
     MetaLearner,
     Params,
 )
+from metalearners.rlearner import OUTCOME_MODEL, RLearner
 from metalearners.tlearner import TLearner
 from metalearners.utils import metalearner_factory, simplify_output
-from metalearners.xlearner import XLearner
+from metalearners.xlearner import CONTROL_EFFECT_MODEL, TREATMENT_EFFECT_MODEL, XLearner
 
 # Chosen arbitrarily.
 _OOS_REFERENCE_VALUE_TOLERANCE = 0.05
@@ -911,3 +914,50 @@ def test_model_reusage(outcome_kind, request):
         xlearner.predict_nuisance(covariates, PROPENSITY_MODEL, 0, False),
         drlearner.predict_nuisance(covariates, PROPENSITY_MODEL, 0, False),
     )
+
+
+@pytest.mark.parametrize(
+    "metalearner_factory, feature_set",
+    [
+        (TLearner, {VARIANT_OUTCOME_MODEL: [0, 1]}),
+        (
+            XLearner,
+            {
+                VARIANT_OUTCOME_MODEL: [0],
+                PROPENSITY_MODEL: [2, 3],
+                TREATMENT_EFFECT_MODEL: [4],
+                CONTROL_EFFECT_MODEL: None,
+            },
+        ),
+        (
+            RLearner,
+            {OUTCOME_MODEL: None, PROPENSITY_MODEL: [4], TREATMENT_MODEL: [3]},
+        ),
+        (
+            DRLearner,
+            {VARIANT_OUTCOME_MODEL: [], PROPENSITY_MODEL: None, TREATMENT_MODEL: [0]},
+        ),
+    ],
+)
+@pytest.mark.parametrize("use_pandas", [False, True])
+def test_evaluate_feature_set_smoke(metalearner_factory, feature_set, rng, use_pandas):
+    n_samples = 100
+    X = rng.standard_normal((n_samples, 5))
+    y = rng.standard_normal(n_samples)
+    w = rng.integers(0, 2, n_samples)
+    if use_pandas:
+        X = pd.DataFrame(X)
+        y = pd.Series(y)
+        w = pd.Series(w)
+
+    ml = metalearner_factory(
+        n_variants=2,
+        is_classification=False,
+        nuisance_model_factory=LinearRegression,
+        treatment_model_factory=LinearRegression,
+        propensity_model_factory=LogisticRegression,
+        feature_set=feature_set,
+        n_folds=2,
+    )
+    ml.fit(X, y, w)
+    ml.evaluate(X, y, w, False)