Parametrize evaluate

CHANGELOG.rst

@@ -15,6 +15,10 @@ Changelog
 0.4.0 (2024-06-18)
 ------------------
 
+* Added ``scoring`` parameter to :meth:`metalearners.metalearner.MetaLearner.evaluate` and
+  implemented the abstract method for the :class:`metalearners.XLearner` and
+  :class:`metalearners.DRLearner`.
+
 * Implemented :meth:`metalearners.cross_fit_estimator.CrossFitEstimator.clone`.
 
 * Added ``n_jobs_base_learners`` to :meth:`metalearners.metalearner.MetaLearner.fit`.

metalearners/_utils.py

@@ -463,3 +463,20 @@ def simplify_output_2d(tensor: np.ndarray) -> np.ndarray:
             "This function requires a regression or a classification with binary outcome "
             "task."
         )
+
+
+# Taken from https://stackoverflow.com/questions/13741998/is-there-a-way-to-let-classes-inherit-the-documentation-of-their-superclass-with
+def copydoc(fromfunc, sep="\n"):
+    """
+    Decorator: Copy the docstring of `fromfunc`
+    """
+
+    def _decorator(func):
+        sourcedoc = fromfunc.__doc__
+        if func.__doc__ is None:
+            func.__doc__ = sourcedoc
+        else:
+            func.__doc__ = sep.join([sourcedoc, func.__doc__])
+        return func
+
+    return _decorator

metalearners/cross_fit_estimator.py

@@ -362,12 +362,17 @@ def __init__(
         self.original_predict_proba = model.predict_proba
 
     def __enter__(self):
-        self.model.predict = partial(  # type: ignore
+        new_predict = partial(
             self.model.predict, is_oos=self.is_oos, oos_method=self.oos_method
         )
-        self.model.predict_proba = partial(  # type: ignore
+        new_predict.__name__ = "predict"  # type: ignore
+        self.model.predict = new_predict  # type: ignore
+
+        new_predict_proba = partial(
             self.model.predict_proba, is_oos=self.is_oos, oos_method=self.oos_method
         )
+        new_predict_proba.__name__ = "predict_proba"  # type: ignore
+        self.model.predict_proba = new_predict_proba  # type: ignore
         return self.model
 
     def __exit__(self, *args):

metalearners/drlearner.py

@@ -1,6 +1,8 @@
 # # Copyright (c) QuantCo 2024-2024
 # # SPDX-License-Identifier: BSD-3-Clause
 
+from collections.abc import Callable, Mapping
+
 import numpy as np
 from joblib import Parallel, delayed
 from typing_extensions import Self
@@ -23,6 +25,7 @@
     VARIANT_OUTCOME_MODEL,
     MetaLearner,
     _ConditionalAverageOutcomeMetaLearner,
+    _evaluate_model,
     _fit_cross_fit_estimator_joblib,
     _ModelSpecifications,
     _ParallelJoblibSpecification,
@@ -148,6 +151,7 @@ def fit(
                 w=w,
                 y=y,
                 treatment_variant=treatment_variant,
+                is_oos=False,
             )
 
             treatment_jobs.append(
@@ -205,37 +209,90 @@ def evaluate(
         w: Vector,
         is_oos: bool,
         oos_method: OosMethod = OVERALL,
-    ) -> dict[str, float | int]:
-        raise NotImplementedError(
-            "This feature is not yet implemented for the DR-Learner."
+        scoring: Mapping[str, list[str | Callable]] | None = None,
+    ) -> dict[str, float]:
+        if scoring is None:
+            scoring = {}
+        self._validate_scoring(scoring=scoring)
+
+        default_metric = (
+            "neg_log_loss" if self.is_classification else "neg_root_mean_squared_error"
+        )
+        masks = []
+        for tv in range(self.n_variants):
+            masks.append(w == tv)
+        variant_outcome_evaluation = _evaluate_model(
+            cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
+            X=[X[w == tv] for tv in range(self.n_variants)],
+            y=[y[w == tv] for tv in range(self.n_variants)],
+            scorers=scoring.get(VARIANT_OUTCOME_MODEL, [default_metric]),
+            model_kind=VARIANT_OUTCOME_MODEL,
+            is_oos=is_oos,
+            oos_method=oos_method,
+            is_treatment=False,
+        )
+
+        propensity_evaluation = _evaluate_model(
+            cfes=self._nuisance_models[PROPENSITY_MODEL],
+            X=[X],
+            y=[w],
+            scorers=scoring.get(PROPENSITY_MODEL, ["neg_log_loss"]),
+            model_kind=PROPENSITY_MODEL,
+            is_oos=is_oos,
+            oos_method=oos_method,
+            is_treatment=False,
         )
 
+        pseudo_outcome: list[np.ndarray] = []
+        for treatment_variant in range(1, self.n_variants):
+            tv_pseudo_outcome = self._pseudo_outcome(
+                X=X,
+                y=y,
+                w=w,
+                treatment_variant=treatment_variant,
+                is_oos=is_oos,
+                oos_method=oos_method,
+            )
+            pseudo_outcome.append(tv_pseudo_outcome)
+
+        treatment_evaluation = _evaluate_model(
+            self._treatment_models[TREATMENT_MODEL],
+            X=[X for _ in range(1, self.n_variants)],
+            y=pseudo_outcome,
+            scorers=scoring.get(TREATMENT_MODEL, ["neg_root_mean_squared_error"]),
+            model_kind=TREATMENT_MODEL,
+            is_oos=is_oos,
+            oos_method=oos_method,
+            is_treatment=True,
+        )
+
+        return variant_outcome_evaluation | propensity_evaluation | treatment_evaluation
+
     def _pseudo_outcome(
         self,
         X: Matrix,
         y: Vector,
         w: Vector,
         treatment_variant: int,
+        is_oos: bool,
+        oos_method: OosMethod = OVERALL,
         epsilon: float = _EPSILON,
     ) -> np.ndarray:
-        """Compute the DR-Learner pseudo outcome.
-
-        Importantly, this method assumes to be applied on in-sample data.
-        In other words, ``is_oos`` will always be set to ``False`` when calling
-        ``predict_nuisance``.
-        """
+        """Compute the DR-Learner pseudo outcome."""
         validate_valid_treatment_variant_not_control(treatment_variant, self.n_variants)
 
         conditional_average_outcome_estimates = (
             self.predict_conditional_average_outcomes(
                 X=X,
-                is_oos=False,
+                is_oos=is_oos,
+                oos_method=oos_method,
             )
         )
 
         propensity_estimates = self.predict_nuisance(
             X=X,
-            is_oos=False,
+            is_oos=is_oos,
+            oos_method=oos_method,
             model_kind=PROPENSITY_MODEL,
             model_ord=0,
         )

metalearners/metalearner.py

@@ -2,14 +2,15 @@
 # # SPDX-License-Identifier: BSD-3-Clause
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Collection
+from collections.abc import Callable, Collection, Mapping, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
 from typing import TypedDict
 
 import numpy as np
 import pandas as pd
 import shap
+from sklearn.metrics import get_scorer
 from sklearn.model_selection import KFold
 from typing_extensions import Self
 
@@ -32,6 +33,7 @@
 from metalearners.cross_fit_estimator import (
     OVERALL,
     CrossFitEstimator,
+    _PredictContext,
 )
 from metalearners.explainer import Explainer
 
@@ -133,6 +135,41 @@ def _validate_n_folds_synchronize(n_folds: dict[str, int]) -> None:
         raise ValueError("Need at least two folds to use synchronization.")
 
 
+def _evaluate_model(
+    cfes: Sequence[CrossFitEstimator],
+    X: Sequence[Matrix],
+    y: Sequence[Vector],
+    scorers: Sequence[str | Callable],
+    model_kind: str,
+    is_oos: bool,
+    is_treatment: bool,
+    oos_method: OosMethod = OVERALL,
+) -> dict[str, float]:
+    """Helper function to evaluate all the models of the same model kind."""
+    prefix = f"{model_kind}_"
+    evaluation_metrics: dict[str, float] = {}
+    for idx, scorer in enumerate(scorers):
+        if isinstance(scorer, str):
+            scorer_str = scorer
+            scorer_call: Callable = get_scorer(scorer)
+        else:
+            scorer_str = f"custom_scorer_{idx}"
+            scorer_call = scorer
+        for i, cfe in enumerate(cfes):
+            if is_treatment:
+                treatment_variant = i + 1
+                index_str = f"{treatment_variant}_vs_0_"
+            else:
+                if len(cfes) == 1:
+                    index_str = ""
+                else:
+                    index_str = f"{i}_"
+            name = f"{prefix}{index_str}{scorer_str}"
+            with _PredictContext(cfe, is_oos, oos_method) as modified_cfe:
+                evaluation_metrics[name] = scorer_call(modified_cfe, X[i], y[i])
+    return evaluation_metrics
+
+
 class _ModelSpecifications(TypedDict):
     # The quotes on MetaLearner are necessary for type hinting as it's not yet defined
     # here. Check https://stackoverflow.com/questions/55320236/does-python-evaluate-type-hinting-of-a-forward-reference
@@ -311,6 +348,16 @@ def _validate_models(self) -> None:
                 factory, predict_method, name=f"treatment model {model_kind}"
             )
 
+    @classmethod
+    def _validate_scoring(cls, scoring: Mapping[str, list[str | Callable]]):
+        if not set(scoring.keys()) <= (
+            set(cls.nuisance_model_specifications().keys())
+            | set(cls.treatment_model_specifications().keys())
+        ):
+            raise ValueError(
+                "scoring dict keys need to be a subset of the model names in the MetaLearner"
+            )
+
     def _qualified_fit_params(
         self,
         fit_params: None | dict,
@@ -824,8 +871,39 @@ def evaluate(
         w: Vector,
         is_oos: bool,
         oos_method: OosMethod = OVERALL,
-    ) -> dict[str, float | int]:
-        """Evaluate all models contained in a MetaLearner."""
+        scoring: Mapping[str, list[str | Callable]] | None = None,
+    ) -> dict[str, float]:
+        r"""Evaluate the models models contained in the MetaLearner.
+
+        ``scoring`` keys must be a subset of the names of the models contained in the
+        MetaLearner, for information about this names check :meth:`~metalearners.metalearner.MetaLearner.nuisance_model_specifications`
+        and :meth:`~metalearners.metalearner.MetaLearner.treatment_model_specifications`.
+        The values must be a list of:
+
+        * ``string`` representing a ``sklearn`` scoring method. Check
+          `here <https://scikit-learn.org/stable/modules/model_evaluation.html#common-cases-predefined-values>`__
+          for the possible values.
+        * ``Callable`` with signature ``scorer(estimator, X, y_true, **kwargs)``. We recommend
+          using `sklearn.metrics.make_scorer <https://scikit-learn.org/stable/modules/generated/sklearn.metrics.make_scorer.html>`_
+          to create this callables.
+
+        If some model name is not present in the keys of ``scoring`` then the default used
+        metrics will be ``neg_log_loss`` if it is a classifier and ``neg_root_mean_squared_error``
+        if it is a regressor.
+
+        The returned dictionary keys have the following structure:
+
+        * For nuisance models:
+
+            * If the cardinality is one:  ``f"{model_kind}_{scorer}"``
+            * If there is one model for each treatment variant (including control):
+              ``f"{model_kind}_{treatment_variant}_{scorer}"``
+
+        * For treatment models: ``f"{model_kind}_{treatment_variant}_vs_0_{scorer}"``
+
+        Where ``scorer`` is the name of the scorer if it is a string and ``"custom_scorer_{idx}"``
+        if it is a callable where ``idx`` is the index in the ``scorers`` list.
+        """
         ...
 
     def explainer(