Merge branch 'parametrize_evaluate' into implement_grid_search

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

@@ -477,7 +477,7 @@ def simplify_output_2d(tensor: np.ndarray) -> np.ndarray:
 # 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`
+    Decorator: Copy the docstring of ``fromfunc``
     """
 
     def _decorator(func):
@@ -489,3 +489,9 @@ def _decorator(func):
         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/drlearner.py

@@ -1,13 +1,12 @@
 # # 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
 
-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,
@@ -209,37 +208,30 @@ def evaluate(
         w: Vector,
         is_oos: bool,
         oos_method: OosMethod = OVERALL,
-        scoring: Mapping[str, list[str | Callable]] | None = None,
+        scoring: Scoring | None = None,
     ) -> dict[str, float]:
-        if scoring is None:
-            scoring = {}
+        safe_scoring = self._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_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=scoring.get(VARIANT_OUTCOME_MODEL, [default_metric]),
+            scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
             model_kind=VARIANT_OUTCOME_MODEL,
             is_oos=is_oos,
             oos_method=oos_method,
-            is_treatment=False,
+            is_treatment_model=False,
         )
 
         propensity_evaluation = _evaluate_model_kind(
             cfes=self._nuisance_models[PROPENSITY_MODEL],
             Xs=[X],
             ys=[w],
-            scorers=scoring.get(PROPENSITY_MODEL, ["neg_log_loss"]),
+            scorers=safe_scoring[PROPENSITY_MODEL],
             model_kind=PROPENSITY_MODEL,
             is_oos=is_oos,
             oos_method=oos_method,
-            is_treatment=False,
+            is_treatment_model=False,
         )
 
         pseudo_outcome: list[np.ndarray] = []
@@ -258,11 +250,11 @@ def evaluate(
             self._treatment_models[TREATMENT_MODEL],
             Xs=[X for _ in range(1, self.n_variants)],
             ys=pseudo_outcome,
-            scorers=scoring.get(TREATMENT_MODEL, ["neg_root_mean_squared_error"]),
+            scorers=safe_scoring[TREATMENT_MODEL],
             model_kind=TREATMENT_MODEL,
             is_oos=is_oos,
             oos_method=oos_method,
-            is_treatment=True,
+            is_treatment_model=True,
         )
 
         return variant_outcome_evaluation | propensity_evaluation | treatment_evaluation

metalearners/metalearner.py

@@ -21,11 +21,13 @@
     OosMethod,
     Params,
     PredictMethod,
+    Scoring,
     SplitIndices,
     Vector,
     _ScikitModel,
 )
 from metalearners._utils import (
+    default_metric,
     index_matrix,
     validate_model_and_predict_method,
     validate_number_positive,
@@ -142,7 +144,7 @@ def _evaluate_model_kind(
     scorers: Sequence[str | Callable],
     model_kind: str,
     is_oos: bool,
-    is_treatment: bool,
+    is_treatment_model: bool,
     oos_method: OosMethod = OVERALL,
 ) -> dict[str, float]:
     """Helper function to evaluate all the models of the same model kind."""
@@ -156,7 +158,7 @@ def _evaluate_model_kind(
             scorer_name = f"custom_scorer_{idx}"
             scorer_callable = scorer
         for i, cfe in enumerate(cfes):
-            if is_treatment:
+            if is_treatment_model:
                 treatment_variant = i + 1
                 index_str = f"{treatment_variant}_vs_0_"
             else:
@@ -877,7 +879,7 @@ def evaluate(
           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.
+          to create such a ``Callable``.
 
         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``
@@ -1025,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):
 

metalearners/rlearner.py

@@ -1,14 +1,13 @@
 # # 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 sklearn.metrics import root_mean_squared_error
 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,
     copydoc,
@@ -335,37 +334,33 @@ def evaluate(
         w: Vector,
         is_oos: bool,
         oos_method: OosMethod = OVERALL,
-        scoring: Mapping[str, list[str | Callable]] | None = None,
+        scoring: Scoring | None = None,
     ) -> dict[str, float]:
         """In the RLearner case, the ``"treatment_model"`` is always evaluated with the
         :func:`~metalearners.rlearner.r_loss` and the ``scoring["treatment_model"]``
         parameter is ignored."""
-        if scoring is None:
-            scoring = {}
+        safe_scoring = self._scoring(scoring)
 
         propensity_evaluation = _evaluate_model_kind(
             cfes=self._nuisance_models[PROPENSITY_MODEL],
             Xs=[X],
             ys=[w],
-            scorers=scoring.get(PROPENSITY_MODEL, ["neg_log_loss"]),
+            scorers=safe_scoring[PROPENSITY_MODEL],
             model_kind=PROPENSITY_MODEL,
             is_oos=is_oos,
             oos_method=oos_method,
-            is_treatment=False,
+            is_treatment_model=False,
         )
 
-        default_metric = (
-            "neg_log_loss" if self.is_classification else "neg_root_mean_squared_error"
-        )
         outcome_evaluation = _evaluate_model_kind(
             cfes=self._nuisance_models[OUTCOME_MODEL],
             Xs=[X],
             ys=[y],
-            scorers=scoring.get(OUTCOME_MODEL, [default_metric]),
+            scorers=safe_scoring[OUTCOME_MODEL],
             model_kind=OUTCOME_MODEL,
             is_oos=is_oos,
             oos_method=oos_method,
-            is_treatment=False,
+            is_treatment_model=False,
         )
 
         # TODO: improve this? generalize it to other metalearners?

metalearners/slearner.py

@@ -2,13 +2,12 @@
 # # SPDX-License-Identifier: BSD-3-Clause
 
 import warnings
-from collections.abc import Callable, Mapping
 
 import numpy as np
 import pandas as pd
 from typing_extensions import Self
 
-from metalearners._typing import Matrix, OosMethod, Vector
+from metalearners._typing import Matrix, OosMethod, Scoring, Vector
 from metalearners._utils import (
     convert_treatment,
     get_one,
@@ -158,14 +157,9 @@ def evaluate(
         w: Vector,
         is_oos: bool,
         oos_method: OosMethod = OVERALL,
-        scoring: Mapping[str, list[str | Callable]] | None = None,
+        scoring: Scoring | None = None,
     ) -> dict[str, float]:
-        if scoring is None:
-            scoring = {}
-
-        default_metric = (
-            "neg_log_loss" if self.is_classification else "neg_root_mean_squared_error"
-        )
+        safe_scoring = self._scoring(scoring)
 
         X_with_w = _append_treatment_to_covariates(
             X, w, self._supports_categoricals, self.n_variants
@@ -174,11 +168,11 @@ def evaluate(
             cfes=self._nuisance_models[_BASE_MODEL],
             Xs=[X_with_w],
             ys=[y],
-            scorers=scoring.get(_BASE_MODEL, [default_metric]),
+            scorers=safe_scoring[_BASE_MODEL],
             model_kind=_BASE_MODEL,
             is_oos=is_oos,
             oos_method=oos_method,
-            is_treatment=False,
+            is_treatment_model=False,
         )
 
     def predict_conditional_average_outcomes(

metalearners/tlearner.py

@@ -2,13 +2,11 @@
 # # 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
 
-from metalearners._typing import Matrix, OosMethod, Vector
+from metalearners._typing import Matrix, OosMethod, Scoring, Vector
 from metalearners._utils import index_matrix
 from metalearners.cross_fit_estimator import OVERALL
 from metalearners.metalearner import (
@@ -116,25 +114,17 @@ def evaluate(
         w: Vector,
         is_oos: bool,
         oos_method: OosMethod = OVERALL,
-        scoring: Mapping[str, list[str | Callable]] | None = None,
+        scoring: Scoring | None = None,
     ) -> dict[str, float]:
-        if scoring is None:
-            scoring = {}
-
-        default_metric = (
-            "neg_log_loss" if self.is_classification else "neg_root_mean_squared_error"
-        )
+        safe_scoring = self._scoring(scoring)
 
-        masks = []
-        for tv in range(self.n_variants):
-            masks.append(w == tv)
         return _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=scoring.get(VARIANT_OUTCOME_MODEL, [default_metric]),
+            scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
             model_kind=VARIANT_OUTCOME_MODEL,
             is_oos=is_oos,
             oos_method=oos_method,
-            is_treatment=False,
+            is_treatment_model=False,
         )