Parametrize evaluate

CHANGELOG.rst

@@ -30,6 +30,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/_typing.py

@@ -1,7 +1,7 @@
 # # Copyright (c) QuantCo 2024-2024
 # # SPDX-License-Identifier: BSD-3-Clause
 
-from collections.abc import Collection, Mapping
+from collections.abc import Callable, Collection, Mapping, Sequence
 from typing import Literal, Protocol, Union
 
 import numpy as np
@@ -29,7 +29,6 @@ class _ScikitModel(Protocol):
 
     # https://stackoverflow.com/questions/54868698/what-type-is-a-sklearn-model/60542986#60542986
     def fit(self, X, y, *params, **kwargs): ...
-
     def predict(self, X, *params, **kwargs): ...
 
     def score(self, X, y, **kwargs): ...
@@ -44,3 +43,7 @@ def set_params(self, **params): ...
 # For instance, if converting the Generator resulting from a call to
 # sklearn.model_selection.KFold.split to a list we obtain this type.
 SplitIndices = list[tuple[np.ndarray, np.ndarray]]
+
+Scorer = str | Callable
+Scorers = Sequence[Scorer]
+Scoring = Mapping[str, Scorers]

metalearners/_utils.py

@@ -463,3 +463,26 @@
             "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
+
+
+def default_metric(predict_method: PredictMethod) -> str:
+    if predict_method == _PREDICT_PROBA:
+        return "neg_log_loss"
+    return "neg_root_mean_squared_error"

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,11 +1,12 @@
 # # Copyright (c) QuantCo 2024-2024
 # # SPDX-License-Identifier: BSD-3-Clause
 
+
 import numpy as np
 from joblib import Parallel, delayed
 from typing_extensions import Self
 
-from metalearners._typing import Matrix, OosMethod, Vector
+from metalearners._typing import Matrix, OosMethod, Scoring, Vector
 from metalearners._utils import (
     clip_element_absolute_value_to_epsilon,
     get_one,
@@ -23,6 +24,7 @@
     VARIANT_OUTCOME_MODEL,
     MetaLearner,
     _ConditionalAverageOutcomeMetaLearner,
+    _evaluate_model_kind,
     _fit_cross_fit_estimator_joblib,
     _ModelSpecifications,
     _ParallelJoblibSpecification,
@@ -148,6 +150,7 @@ def fit(
                 w=w,
                 y=y,
                 treatment_variant=treatment_variant,
+                is_oos=False,
             )
 
             treatment_jobs.append(
@@ -205,37 +208,85 @@ 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: Scoring | None = None,
+    ) -> dict[str, float]:
+        safe_scoring = self._scoring(scoring)
+
+        masks = []
+        for tv in range(self.n_variants):
+            masks.append(w == tv)
+        variant_outcome_evaluation = _evaluate_model_kind(
+            cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
+            Xs=[X[w == tv] for tv in range(self.n_variants)],
+            ys=[y[w == tv] for tv in range(self.n_variants)],
+            scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
+            model_kind=VARIANT_OUTCOME_MODEL,
+            is_oos=is_oos,
+            oos_method=oos_method,
+            is_treatment=False,
+        )
+
+        propensity_evaluation = _evaluate_model_kind(
+            cfes=self._nuisance_models[PROPENSITY_MODEL],
+            Xs=[X],
+            ys=[w],
+            scorers=safe_scoring[PROPENSITY_MODEL],
+            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_kind(
+            self._treatment_models[TREATMENT_MODEL],
+            Xs=[X for _ in range(1, self.n_variants)],
+            ys=pseudo_outcome,
+            scorers=safe_scoring[TREATMENT_MODEL],
+            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
 
@@ -20,18 +21,21 @@
     OosMethod,
     Params,
     PredictMethod,
+    Scoring,
     SplitIndices,
     Vector,
     _ScikitModel,
 )
 from metalearners._utils import (
+    default_metric,
     index_matrix,
     validate_model_and_predict_method,
     validate_number_positive,
 )
 from metalearners.cross_fit_estimator import (
     OVERALL,
     CrossFitEstimator,
+    _PredictContext,
 )
 from metalearners.explainer import Explainer
 
@@ -133,6 +137,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_kind(
+    cfes: Sequence[CrossFitEstimator],
+    Xs: Sequence[Matrix],
+    ys: 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_name = scorer
+            scorer_callable: Callable = get_scorer(scorer)
+        else:
+            scorer_name = f"custom_scorer_{idx}"
+            scorer_callable = 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_name}"
+            with _PredictContext(cfe, is_oos, oos_method) as modified_cfe:
+                evaluation_metrics[name] = scorer_callable(modified_cfe, Xs[i], ys[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
@@ -825,8 +864,40 @@ 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 MetaLearner.
+
+        The keys in ``scoring`` which are not a name of a model contained in the MetaLearner
+        will be ignored, 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(
@@ -956,6 +1027,27 @@ def shap_values(
             shap_explainer_params=shap_explainer_params,
         )
 
+    def _scoring(self, scoring: Scoring | None) -> Scoring:
+
+        def _default_scoring() -> Scoring:
+            return {
+                nuisance_model: [
+                    default_metric(model_specifications["predict_method"](self))
+                ]
+                for nuisance_model, model_specifications in self.nuisance_model_specifications().items()
+            } | {
+                treatment_model: [
+                    default_metric(model_specifications["predict_method"](self))
+                ]
+                for treatment_model, model_specifications in self.treatment_model_specifications().items()
+            }
+
+        default_scoring = _default_scoring()
+
+        if scoring is None:
+            return default_scoring
+        return dict(default_scoring) | dict(scoring)
+
 
 class _ConditionalAverageOutcomeMetaLearner(MetaLearner, ABC):