diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d04fad1f..50f2ee47 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -7,6 +7,28 @@ Changelog ========= +0.8.0 (2024-07-22) +------------------ + +**New features** + +* Added :meth:`metalearners.metalearner.MetaLearner.fit_all_nuisance` and + :meth:`metalearners.metalearner.MetaLearner.fit_all_treatment`. + +* Add optional ``store_raw_results`` and ``store_results`` parameters to :class:`metalearners.grid_search.MetaLearnerGridSearch`. + +* Renamed :class:`metalearners.grid_search._GSResult` to :class:`metalearners.grid_search.GSResult`. + +* Added ``grid_size_`` attribute to :class:`metalearners.grid_search.MetaLearnerGridSearch`. + +* 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) ------------------ diff --git a/docs/examples/example_gridsearch.ipynb b/docs/examples/example_gridsearch.ipynb index 586c5d02..ada0a007 100644 --- a/docs/examples/example_gridsearch.ipynb +++ b/docs/examples/example_gridsearch.ipynb @@ -327,6 +327,26 @@ "gs.results_" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "What if I run out of memory?\n", + "----------------------------\n", + "\n", + "If you're conducting an optimization task over a large grid with a substantial dataset,\n", + "it is possible that memory usage issues may arise. To try to solve these, you can minimize\n", + "memory usage by adjusting your settings.\n", + "\n", + "In that case you can set ``store_raw_results=False``, the grid search will then operate\n", + "with a generator rather than a list, significantly reducing memory usage.\n", + "\n", + "If the ``results_ DataFrame`` is what you're after, you can simply set ``store_results=True``.\n", + "However, if you aim to iterate over the {class}`~metalearners.metalearner.MetaLearner` objects,\n", + "you can set ``store_results=False``. Consequently, ``raw_results_`` will become a generator\n", + "object yielding {class}`~metalearners.grid_search.GSResult`." + ] + }, { "cell_type": "markdown", "metadata": {}, diff --git a/metalearners/_typing.py b/metalearners/_typing.py index 76b31ba5..95b66b8b 100644 --- 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/cross_fit_estimator.py b/metalearners/cross_fit_estimator.py index 5b141ac6..aa112c03 100644 --- a/metalearners/cross_fit_estimator.py +++ b/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() diff --git a/metalearners/drlearner.py b/metalearners/drlearner.py index c0bc92ba..a7df898d 100644 --- a/metalearners/drlearner.py +++ b/metalearners/drlearner.py @@ -13,6 +13,7 @@ OosMethod, Params, Scoring, + SplitIndices, Vector, _ScikitModel, ) @@ -128,7 +129,7 @@ def __init__( ) self.adaptive_clipping = adaptive_clipping - def fit( + def fit_all_nuisance( self, X: Matrix, y: Vector, @@ -148,10 +149,12 @@ def fit( for treatment_variant in range(self.n_variants): self._treatment_variants_indices.append(w == treatment_variant) + self._cv_split_indices: SplitIndices | None + if synchronize_cross_fitting: - cv_split_indices = self._split(X) + self._cv_split_indices = self._split(X) else: - cv_split_indices = None + self._cv_split_indices = None nuisance_jobs: list[_ParallelJoblibSpecification | None] = [] for treatment_variant in range(self.n_variants): @@ -176,7 +179,7 @@ def fit( model_ord=0, n_jobs_cross_fitting=n_jobs_cross_fitting, fit_params=qualified_fit_params[NUISANCE][PROPENSITY_MODEL], - cv=cv_split_indices, + cv=self._cv_split_indices, ) ) @@ -189,6 +192,25 @@ def fit( self._assign_joblib_nuisance_results(results) + return self + + def fit_all_treatment( + self, + X: Matrix, + y: Vector, + w: Vector, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + ) -> Self: + if not hasattr(self, "_cv_split_indices"): + raise ValueError( + "The nuisance models need to be fitted before fitting the treatment models." + "In particular, the MetaLearner's attribute _cv_split_indices, " + "typically set during nuisance fitting, does not exist." + ) + qualified_fit_params = self._qualified_fit_params(fit_params) treatment_jobs: list[_ParallelJoblibSpecification] = [] for treatment_variant in range(1, self.n_variants): pseudo_outcomes = self._pseudo_outcome( @@ -207,9 +229,10 @@ def fit( model_ord=treatment_variant - 1, n_jobs_cross_fitting=n_jobs_cross_fitting, fit_params=qualified_fit_params[TREATMENT][TREATMENT_MODEL], - cv=cv_split_indices, + cv=self._cv_split_indices, ) ) + parallel = Parallel(n_jobs=n_jobs_base_learners) results = parallel( delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs ) @@ -267,6 +290,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 +302,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 +326,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/grid_search.py b/metalearners/grid_search.py index cc9c732e..93f19818 100644 --- a/metalearners/grid_search.py +++ b/metalearners/grid_search.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause import time -from collections.abc import Mapping, Sequence +from collections.abc import Generator, Mapping, Sequence from dataclasses import dataclass from typing import Any @@ -17,7 +17,8 @@ @dataclass(frozen=True) class _FitAndScoreJob: - metalearner: MetaLearner + metalearner_factory: type[MetaLearner] + metalearner_params: dict[str, Any] X_train: Matrix y_train: Vector w_train: Vector @@ -32,7 +33,7 @@ class _FitAndScoreJob: @dataclass(frozen=True) -class _GSResult: +class GSResult: r"""Result from a single grid search evaluation.""" metalearner: MetaLearner @@ -42,15 +43,14 @@ class _GSResult: score_time: float -def _fit_and_score(job: _FitAndScoreJob) -> _GSResult: +def _fit_and_score(job: _FitAndScoreJob) -> GSResult: start_time = time.time() - job.metalearner.fit( - job.X_train, job.y_train, job.w_train, **job.metalerner_fit_params - ) + ml = job.metalearner_factory(**job.metalearner_params) + ml.fit(job.X_train, job.y_train, job.w_train, **job.metalerner_fit_params) fit_time = time.time() - start_time start_time = time.time() - train_scores = job.metalearner.evaluate( + train_scores = ml.evaluate( X=job.X_train, y=job.y_train, w=job.w_train, @@ -58,7 +58,7 @@ def _fit_and_score(job: _FitAndScoreJob) -> _GSResult: scoring=job.scoring, ) if job.X_test is not None and job.y_test is not None and job.w_test is not None: - test_scores = job.metalearner.evaluate( + test_scores = ml.evaluate( X=job.X_test, y=job.y_test, w=job.w_test, @@ -69,8 +69,8 @@ def _fit_and_score(job: _FitAndScoreJob) -> _GSResult: else: test_scores = None score_time = time.time() - start_time - return _GSResult( - metalearner=job.metalearner, + return GSResult( + metalearner=ml, fit_time=fit_time, score_time=score_time, train_scores=train_scores, @@ -78,7 +78,9 @@ def _fit_and_score(job: _FitAndScoreJob) -> _GSResult: ) -def _format_results(results: Sequence[_GSResult]) -> pd.DataFrame: +def _format_results( + results: list[GSResult] | Generator[GSResult, None, None] +) -> pd.DataFrame: rows = [] for result in results: row: dict[str, str | int | float] = {} @@ -180,11 +182,33 @@ class MetaLearnerGridSearch: ``verbose`` will be passed to `joblib.Parallel `_. - After fitting a dataframe with the results will be available in `results_`. + ``store_raw_results`` and ``store_results`` define which and how the results are saved + after calling :meth:`~metalearners.grid_search.MetaLearnerGridSearch.fit` depending on + their values: + + * Both are ``True`` (default): ``raw_results_`` will be a list of + :class:`~metalearners.grid_search.GSResult` with all the results and ``results_`` + will be a DataFrame with the processed results. + * ``store_raw_results=True`` and ``store_results=False``: ``raw_results_`` will be a + list of :class:`~metalearners.grid_search.GSResult` with all the results + and ``results`` will be ``None``. + * ``store_raw_results=False`` and ``store_results=True``: ``raw_results_`` will be + ``None`` and ``results_`` will be a DataFrame with the processed results. + * Both are ``False``: ``raw_results_`` will be a generator which yields a + :class:`~metalearners.grid_search.GSResult` for each configuration and ``results`` + will be None. This configuration can be useful in the case the grid search is big + and you do not want to store all MetaLearners objects rather evaluate them after + fitting each one and just store one. + + ``grid_size_`` will contain the number of hyperparameter combinations after fitting. + This attribute may be useful in the case ``store_raw_results = False`` and ``store_results = False``. + In that case, the generator object returned in ``raw_results_`` doesn't trigger the fitting + of individual metalearners until explicitly requested, e.g. in a loop. This attribute + can be use to track the progress, for instance, by creating a progress bar or a similar utility. + + For an illustration see :ref:`our example on Tuning hyperparameters of a MetaLearner with MetaLearnerGridSearch `. """ - # TODO: Add a reference to a docs example once it is written. - def __init__( self, metalearner_factory: type[MetaLearner], @@ -195,6 +219,8 @@ def __init__( n_jobs: int | None = None, random_state: int | None = None, verbose: int = 0, + store_raw_results: bool = True, + store_results: bool = True, ): self.metalearner_factory = metalearner_factory self.metalearner_params = metalearner_params @@ -202,9 +228,8 @@ def __init__( self.n_jobs = n_jobs self.random_state = random_state self.verbose = verbose - - self.raw_results_: Sequence[_GSResult] | None = None - self.results_: pd.DataFrame | None = None + self.store_raw_results = store_raw_results + self.store_results = store_results all_base_models = set( metalearner_factory.nuisance_model_specifications().keys() @@ -286,20 +311,33 @@ def fit( } propensity_model_params = params.get(PROPENSITY_MODEL, None) - ml = self.metalearner_factory( - **self.metalearner_params, - nuisance_model_factory=nuisance_model_factory, - treatment_model_factory=treatment_model_factory, - propensity_model_factory=propensity_model_factory, - nuisance_model_params=nuisance_model_params, - treatment_model_params=treatment_model_params, - propensity_model_params=propensity_model_params, - random_state=self.random_state, - ) + grid_metalearner_params = { + "nuisance_model_factory": nuisance_model_factory, + "treatment_model_factory": treatment_model_factory, + "propensity_model_factory": propensity_model_factory, + "nuisance_model_params": nuisance_model_params, + "treatment_model_params": treatment_model_params, + "propensity_model_params": propensity_model_params, + "random_state": self.random_state, + } + + if ( + len( + shared_keys := set(grid_metalearner_params.keys()) + & set(self.metalearner_params.keys()) + ) + > 0 + ): + raise ValueError( + f"{shared_keys} should not be specified in metalearner_params as " + "they are used internally. Please use the correct parameters." + ) jobs.append( _FitAndScoreJob( - metalearner=ml, + metalearner_factory=self.metalearner_factory, + metalearner_params=dict(self.metalearner_params) + | grid_metalearner_params, X_train=X, y_train=y, w_train=w, @@ -312,7 +350,17 @@ def fit( ) ) - parallel = Parallel(n_jobs=self.n_jobs, verbose=self.verbose) - raw_results = parallel(delayed(_fit_and_score)(job) for job in jobs) - self.raw_results_ = raw_results - self.results_ = _format_results(results=raw_results) + self.grid_size_ = len(jobs) + self.raw_results_: list[GSResult] | Generator[GSResult, None, None] | None + self.results_: pd.DataFrame | None = None + + return_as = "list" if self.store_raw_results else "generator_unordered" + parallel = Parallel( + n_jobs=self.n_jobs, verbose=self.verbose, return_as=return_as + ) + self.raw_results_ = parallel(delayed(_fit_and_score)(job) for job in jobs) + if self.store_results: + self.results_ = _format_results(results=self.raw_results_) # type: ignore + if not self.store_raw_results: + # The generator will be empty so we replace it with None + self.raw_results_ = None diff --git a/metalearners/metalearner.py b/metalearners/metalearner.py index 6214af96..3b2e456a 100644 --- 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 @@ -744,6 +746,41 @@ def _assign_joblib_treatment_results( ] = result.cross_fit_estimator @abstractmethod + def fit_all_nuisance( + self, + X: Matrix, + y: Vector, + w: Vector, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + ) -> Self: + """Fit all nuisance models of the MetaLearner. + + If pre-fitted models were passed at instantiation, these are never refitted. + + For the parameters check :meth:`metalearners.metalearner.MetaLearner.fit`. + """ + ... + + @abstractmethod + def fit_all_treatment( + self, + X: Matrix, + y: Vector, + w: Vector, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + ) -> Self: + """Fit all treatment models of the MetaLearner. + + For the parameters check :meth:`metalearners.metalearner.MetaLearner.fit`. + """ + ... + def fit( self, X: Matrix, @@ -791,7 +828,27 @@ def fit( the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification. """ - ... + self.fit_all_nuisance( + X=X, + y=y, + w=w, + n_jobs_cross_fitting=n_jobs_cross_fitting, + fit_params=fit_params, + synchronize_cross_fitting=synchronize_cross_fitting, + n_jobs_base_learners=n_jobs_base_learners, + ) + + self.fit_all_treatment( + X=X, + y=y, + w=w, + n_jobs_cross_fitting=n_jobs_cross_fitting, + fit_params=fit_params, + synchronize_cross_fitting=synchronize_cross_fitting, + n_jobs_base_learners=n_jobs_base_learners, + ) + + return self def predict_nuisance( self, diff --git a/metalearners/rlearner.py b/metalearners/rlearner.py index b7e9297f..965b928c 100644 --- a/metalearners/rlearner.py +++ b/metalearners/rlearner.py @@ -156,7 +156,7 @@ def _supports_multi_treatment(cls) -> bool: def _supports_multi_class(cls) -> bool: return False - def fit( + def fit_all_nuisance( self, X: Matrix, y: Vector, @@ -165,14 +165,10 @@ def fit( fit_params: dict | None = None, synchronize_cross_fitting: bool = True, n_jobs_base_learners: int | None = None, - epsilon: float = _EPSILON, ) -> Self: - self._validate_treatment(w) self._validate_outcome(y, w) - self._variants_indices = [] - qualified_fit_params = self._qualified_fit_params(fit_params) self._validate_fit_params(qualified_fit_params) @@ -214,7 +210,22 @@ def fit( ) self._assign_joblib_nuisance_results(results) + return self + + def fit_all_treatment( + self, + X: Matrix, + y: Vector, + w: Vector, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + epsilon: float = _EPSILON, + ) -> Self: + qualified_fit_params = self._qualified_fit_params(fit_params) treatment_jobs: list[_ParallelJoblibSpecification] = [] + self._variants_indices = [] for treatment_variant in range(1, self.n_variants): is_treatment = w == treatment_variant @@ -246,6 +257,7 @@ def fit( n_jobs_cross_fitting=n_jobs_cross_fitting, ) ) + parallel = Parallel(n_jobs=n_jobs_base_learners) results = parallel( delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs ) @@ -352,6 +364,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 +376,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 +428,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 index 7158cca2..3e1d8078 100644 --- a/metalearners/slearner.py +++ b/metalearners/slearner.py @@ -141,7 +141,7 @@ def __init__( random_state=random_state, ) - def fit( + def fit_all_nuisance( self, X: Matrix, y: Vector, @@ -175,6 +175,18 @@ def fit( ) return self + def fit_all_treatment( + self, + X: Matrix, + y: Vector, + w: Vector, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + ) -> Self: + return self + def predict( self, X: Matrix, @@ -212,6 +224,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 index cc61fe5c..d512ab12 100644 --- a/metalearners/tlearner.py +++ b/metalearners/tlearner.py @@ -48,7 +48,7 @@ def _supports_multi_treatment(cls) -> bool: def _supports_multi_class(cls) -> bool: return True - def fit( + def fit_all_nuisance( self, X: Matrix, y: Vector, @@ -92,6 +92,18 @@ def fit( self._assign_joblib_nuisance_results(results) return self + def fit_all_treatment( + self, + X: Matrix, + y: Vector, + w: Vector, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + ) -> Self: + return self + def predict( self, X, @@ -126,4 +138,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 index 109619d7..7e8a1b21 100644 --- a/metalearners/xlearner.py +++ b/metalearners/xlearner.py @@ -74,7 +74,7 @@ def _supports_multi_treatment(cls) -> bool: def _supports_multi_class(cls) -> bool: return False - def fit( + def fit_all_nuisance( self, X: Matrix, y: Vector, @@ -91,7 +91,7 @@ def fit( qualified_fit_params = self._qualified_fit_params(fit_params) - cvs: list = [] + self._cvs: list = [] for treatment_variant in range(self.n_variants): self._treatment_variants_indices.append(w == treatment_variant) @@ -101,7 +101,7 @@ def fit( ) else: cv_split_indices = None - cvs.append(cv_split_indices) + self._cvs.append(cv_split_indices) nuisance_jobs: list[_ParallelJoblibSpecification | None] = [] for treatment_variant in range(self.n_variants): @@ -115,7 +115,7 @@ def fit( model_ord=treatment_variant, n_jobs_cross_fitting=n_jobs_cross_fitting, fit_params=qualified_fit_params[NUISANCE][VARIANT_OUTCOME_MODEL], - cv=cvs[treatment_variant], + cv=self._cvs[treatment_variant], ) ) @@ -138,6 +138,32 @@ def fit( ) self._assign_joblib_nuisance_results(results) + return self + + def fit_all_treatment( + self, + X: Matrix, + y: Vector, + w: Vector, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + ) -> Self: + if self._treatment_variants_indices is None: + raise ValueError( + "The nuisance models need to be fitted before fitting the treatment models." + "In particular, the MetaLearner's attribute _treatment_variant_indices, " + "typically set during nuisance fitting, is None." + ) + if not hasattr(self, "_cvs"): + raise ValueError( + "The nuisance models need to be fitted before fitting the treatment models." + "In particular, the MetaLearner's attribute _cvs, " + "typically set during nuisance fitting, does not exist." + ) + qualified_fit_params = self._qualified_fit_params(fit_params) + treatment_jobs: list[_ParallelJoblibSpecification] = [] for treatment_variant in range(1, self.n_variants): imputed_te_control, imputed_te_treatment = self._pseudo_outcome( @@ -153,7 +179,7 @@ def fit( model_ord=treatment_variant - 1, n_jobs_cross_fitting=n_jobs_cross_fitting, fit_params=qualified_fit_params[TREATMENT][TREATMENT_EFFECT_MODEL], - cv=cvs[treatment_variant], + cv=self._cvs[treatment_variant], ) ) @@ -165,10 +191,11 @@ def fit( model_ord=treatment_variant - 1, n_jobs_cross_fitting=n_jobs_cross_fitting, fit_params=qualified_fit_params[TREATMENT][CONTROL_EFFECT_MODEL], - cv=cvs[0], + cv=self._cvs[0], ) ) + parallel = Parallel(n_jobs=n_jobs_base_learners) results = parallel( delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs ) @@ -300,6 +327,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 +339,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 +360,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 +372,7 @@ def evaluate( is_oos=is_oos, oos_method=oos_method, is_treatment_model=True, + feature_set=self.feature_set[CONTROL_EFFECT_MODEL], ) return ( diff --git a/tests/test_cross_fit_estimator.py b/tests/test_cross_fit_estimator.py index 38223994..4165a06c 100644 --- a/tests/test_cross_fit_estimator.py +++ b/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) diff --git a/tests/test_grid_search.py b/tests/test_grid_search.py index e29d3d3b..23e4cbdb 100644 --- a/tests/test_grid_search.py +++ b/tests/test_grid_search.py @@ -2,7 +2,10 @@ # SPDX-License-Identifier: BSD-3-Clause +from types import GeneratorType + import numpy as np +import pandas as pd import pytest from lightgbm import LGBMClassifier, LGBMRegressor from sklearn.linear_model import LinearRegression, LogisticRegression @@ -153,6 +156,7 @@ def test_metalearnergridsearch_smoke( assert gs.results_ is not None assert gs.results_.shape[0] == expected_n_configs assert gs.results_.index.names == expected_index_cols + assert gs.grid_size_ == expected_n_configs train_scores_cols = set( c[6:] for c in list(gs.results_.columns) if c.startswith("train_") @@ -259,3 +263,65 @@ def test_metalearnergridsearch_reuse_propensity_smoke(grid_search_data): assert gs.results_ is not None assert gs.results_.shape[0] == 2 assert len(gs.results_.index.names) == 5 + + +@pytest.mark.parametrize( + "store_raw_results, store_results, expected_type_raw_results, expected_type_results", + [ + (True, True, list, pd.DataFrame), + (True, False, list, type(None)), + (False, True, type(None), pd.DataFrame), + (False, False, GeneratorType, type(None)), + ], +) +def test_metalearnergridsearch_store( + store_raw_results, + store_results, + expected_type_raw_results, + expected_type_results, + grid_search_data, +): + X, _, y, w, X_test, _, y_test, w_test = grid_search_data + n_variants = len(np.unique(w)) + + metalearner_params = { + "is_classification": False, + "n_variants": n_variants, + "n_folds": 2, + } + + gs = MetaLearnerGridSearch( + metalearner_factory=SLearner, + metalearner_params=metalearner_params, + base_learner_grid={"base_model": [LinearRegression, LGBMRegressor]}, + param_grid={"base_model": {"LGBMRegressor": {"n_estimators": [1, 2]}}}, + store_raw_results=store_raw_results, + store_results=store_results, + ) + + gs.fit(X, y, w, X_test, y_test, w_test) + assert isinstance(gs.raw_results_, expected_type_raw_results) + assert isinstance(gs.results_, expected_type_results) + + +def test_metalearnergridsearch_error(grid_search_data): + X, _, y, w, X_test, _, y_test, w_test = grid_search_data + n_variants = len(np.unique(w)) + + metalearner_params = { + "is_classification": False, + "n_variants": n_variants, + "n_folds": 2, + "random_state": 1, + } + + gs = MetaLearnerGridSearch( + metalearner_factory=SLearner, + metalearner_params=metalearner_params, + base_learner_grid={"base_model": [LinearRegression, LGBMRegressor]}, + param_grid={"base_model": {"LGBMRegressor": {"n_estimators": [1, 2]}}}, + ) + with pytest.raises( + ValueError, match="should not be specified in metalearner_params" + ): + gs.fit(X, y, w, X_test, y_test, w_test) diff --git a/tests/test_learner.py b/tests/test_learner.py index 4aa80727..9971f0e0 100644 --- a/tests/test_learner.py +++ b/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) diff --git a/tests/test_metalearner.py b/tests/test_metalearner.py index 371f971c..03633c2e 100644 --- a/tests/test_metalearner.py +++ b/tests/test_metalearner.py @@ -68,7 +68,7 @@ def _supports_multi_class(cls) -> bool: def _validate_models(self) -> None: ... - def fit( + def fit_all_nuisance( self, X, y, @@ -83,6 +83,18 @@ def fit( self.nuisance_model_specifications()[model_kind]["cardinality"](self) ): self.fit_nuisance(X, y, model_kind, model_ord) + return self + + def fit_all_treatment( + self, + X, + y, + w, + n_jobs_cross_fitting: int | None = None, + fit_params: dict | None = None, + synchronize_cross_fitting: bool = True, + n_jobs_base_learners: int | None = None, + ): for model_kind in self.__class__.treatment_model_specifications(): for model_ord in range( self.treatment_model_specifications()[model_kind]["cardinality"](self)