diff --git a/econml/dr/_drlearner.py b/econml/dr/_drlearner.py index bed549b2c..1a0ab9265 100644 --- a/econml/dr/_drlearner.py +++ b/econml/dr/_drlearner.py @@ -246,6 +246,9 @@ def _compute_trim_mask(self, T_pred): def fit(self, Y, T, X=None, W=None, *, nuisances, sample_weight=None, freq_weight=None, sample_var=None, groups=None): Y_pred, T_pred, T_pred_raw = nuisances + if sample_var is not None and np.shape(sample_var) not in [(len(Y),), (len(Y), 1)]: + raise ValueError(f"sample_var must have shape ({len(Y)},) or ({len(Y)}, 1), " + f"instead got {np.shape(sample_var)}.") # Apply sample trimming based on Crump et al. (2009) if enabled # Use raw (unclipped) propensities for trimming decisions @@ -271,7 +274,11 @@ def fit(self, Y, T, X=None, W=None, *, nuisances, self.sensitivity_params = dr_sensitivity_values(Y, T, Y_pred, T_pred) T_complete = np.hstack(((np.all(T == 0, axis=1) * 1).reshape(-1, 1), T)) - propensities = np.sum(T_pred * T_complete, axis=1).reshape((T.shape[0],)) + weighted_sample_var = None + if sample_var is not None: + # Each contrast depends on the observed outcome only for its treatment or control arm. + outcome_derivative = T / T_pred[:, 1:] - T_complete[:, [0]] / T_pred[:, [0]] + weighted_sample_var = sample_var.reshape(-1, 1) * outcome_derivative**2 self.d_y = Y_pred.shape[1:-1] # track whether there's a Y dimension (must be a singleton) self.d_t = Y_pred.shape[-1] - 1 # track # of treatment (exclude baseline treatment) @@ -283,18 +290,18 @@ def fit(self, Y, T, X=None, W=None, *, nuisances, ys = Y_pred[..., 1:] - Y_pred[..., [0]] # subtract control results from each other arm if self.d_y: # need to squeeze out singleton so that we fit on 2D array ys = ys.squeeze(1) - weighted_sample_var = np.tile((sample_var / propensities**2).reshape((-1, 1)), - self.d_t) if sample_var is not None else None filtered_kwargs = filter_none_kwargs(sample_weight=sample_weight, freq_weight=freq_weight, sample_var=weighted_sample_var) self.model_cate = self._model_final.fit(X, ys, **filtered_kwargs) else: - weighted_sample_var = sample_var / propensities**2 if sample_var is not None else None - filtered_kwargs = filter_none_kwargs(sample_weight=sample_weight, - freq_weight=freq_weight, sample_var=weighted_sample_var) - self.models_cate = [clone(self._model_final, safe=False).fit(X, Y_pred[..., t] - Y_pred[..., 0], - **filtered_kwargs) - for t in np.arange(1, Y_pred.shape[-1])] + self.models_cate = [] + for t in range(self.d_t): + target_var = (weighted_sample_var[:, t].reshape((len(Y),) + self.d_y) + if weighted_sample_var is not None else None) + filtered_kwargs = filter_none_kwargs(sample_weight=sample_weight, + freq_weight=freq_weight, sample_var=target_var) + self.models_cate.append(clone(self._model_final, safe=False).fit( + X, Y_pred[..., t + 1] - Y_pred[..., 0], **filtered_kwargs)) return self def predict(self, X=None): @@ -700,7 +707,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam Parameters ---------- - Y: (n,) vector of length n + Y: (n,) or (n, 1) array_like Outcomes for each sample T: (n,) vector of length n Treatments for each sample @@ -714,9 +721,11 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam Weight for the observation. Observation i is treated as the mean outcome of freq_weight[i] independent observations. When ``sample_var`` is not None, this should be provided. - sample_var : (n,) nd array_like, optional + sample_var : (n,) or (n, 1) array_like, optional Variance of the outcome(s) of the original freq_weight[i] observations that were used to - compute the mean outcome represented by observation i. + compute the mean outcome represented by observation i. Either shape is accepted for + vector or singleton-column outcomes. Variances are propagated separately to each + treatment-versus-control contrast. groups: (n,) vector, optional All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer @@ -1341,7 +1350,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam Parameters ---------- - Y: (n,) vector of length n + Y: (n,) or (n, 1) array_like Outcomes for each sample T: (n,) vector of length n Treatments for each sample @@ -1355,9 +1364,11 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam Weight for the observation. Observation i is treated as the mean outcome of freq_weight[i] independent observations. When ``sample_var`` is not None, this should be provided. - sample_var : (n,) nd array_like, optional + sample_var : (n,) or (n, 1) array_like, optional Variance of the outcome(s) of the original freq_weight[i] observations that were used to - compute the mean outcome represented by observation i. + compute the mean outcome represented by observation i. Either shape is accepted for + vector or singleton-column outcomes. Variances are propagated separately to each + treatment-versus-control contrast. groups: (n,) vector, optional All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer diff --git a/econml/inference/_inference.py b/econml/inference/_inference.py index 82fdb19c7..bbaedd4cc 100644 --- a/econml/inference/_inference.py +++ b/econml/inference/_inference.py @@ -1444,6 +1444,9 @@ def conf_int_point(self, *, alpha=None, tol=None): """ Get the confidence interval of the point estimate of each treatment on each outcome for sample X. + An outcome/treatment interval is NaN if any sample mean or standard error is nonfinite, + or any standard error is nonpositive. Other intervals are unaffected. + Parameters ---------- alpha: float in [0, 1], optional @@ -1576,12 +1579,14 @@ def _print(self, *, alpha=None, value=None, decimals=None, tol=None, output_name def _mixture_ppf(self, alpha, mean, stderr, tol): """Get the confidence interval of mixture gaussian distribution.""" - # if stderr is zero, ppf will return nans and the loop below would never terminate - # so bail out early; note that it might be possible to correct the algorithm for - # this scenario, but since scipy's cdf returns nan whenever scale is zero it won't - # be clean - if (np.any(stderr == 0)): - return np.full(shape(mean)[1:], np.nan) + # Invalid normal parameters prevent convergence. Exclude only the affected + # outcome/treatment columns from the solver. + invalid = np.any((stderr <= 0) | ~np.isfinite(stderr) | ~np.isfinite(mean), axis=0) + if np.any(invalid): + result = np.full(shape(mean)[1:], np.nan) + if not np.all(invalid): + result[~invalid] = self._mixture_ppf(alpha, mean[:, ~invalid], stderr[:, ~invalid], tol) + return result mix_ppf = scipy.stats.norm.ppf(alpha, loc=mean, scale=stderr) lower = np.min(mix_ppf, axis=0) upper = np.max(mix_ppf, axis=0) diff --git a/econml/sklearn_extensions/model_selection.py b/econml/sklearn_extensions/model_selection.py index 8f86a3662..422594975 100644 --- a/econml/sklearn_extensions/model_selection.py +++ b/econml/sklearn_extensions/model_selection.py @@ -465,8 +465,15 @@ def _convert_linear_regression(model, new_cls, extra_attrs=["positive"]): def _to_elasticNet(model: ElasticNetCV, args, kwargs, is_lasso=False, cls=None, extra_attrs=[]): - # We need an R^2 score to compare to other models; ElasticNetCV doesn't provide it, - # but we can calculate it ourselves from the MSE plus the variance of the target y + # Retain the fast R^2 approximation: mse_path_ already averages squared errors over + # outputs, so normalize by the mean of the per-output variances. Pooling outputs + # would incorrectly count differences between their means as target variation. + # + # Scores are still not directly comparable to other selectors' default R^2 scores: + # we use variances computed over all supplied observations rather than within each + # validation fold, and aggregate outputs differently. Consistent scoring requires + # broader evaluation changes because mse_path_ no longer contains per-output errors. + # The denominator also remains unweighted when CV uses sample weights. y = signature(model.fit).bind(*args, **kwargs).arguments["y"] cls = cls or (Lasso if is_lasso else ElasticNet) new_model = _convert_linear_regression(model, cls, extra_attrs + ['selection', 'warm_start', 'dual_gap_', @@ -477,7 +484,7 @@ def _to_elasticNet(model: ElasticNetCV, args, kwargs, is_lasso=False, cls=None, _copy_to(model, new_model, ["l1_ratio"], True) # max R^2 corresponds to min MSE min_mse = np.min(np.mean(model.mse_path_, axis=-1)) # last dimension in mse_path is folds, so average over that - r2 = 1 - min_mse / np.var(y) # R^2 = 1 - MSE / Var(y) + r2 = 1 - min_mse / np.mean(np.var(y, axis=0)) # R^2 = 1 - MSE / Var(y) return new_model, r2 diff --git a/econml/tests/test_drlearner.py b/econml/tests/test_drlearner.py index eaced8d87..7386d6cec 100644 --- a/econml/tests/test_drlearner.py +++ b/econml/tests/test_drlearner.py @@ -17,6 +17,7 @@ from sklearn.preprocessing import FunctionTransformer, PolynomialFeatures from econml.dr import DRLearner, LinearDRLearner, SparseLinearDRLearner, ForestDRLearner +from econml.dr._drlearner import _ModelFinal from econml.inference import BootstrapInference, StatsModelsInferenceDiscrete from econml.utilities import get_feature_names_or_default, shape from econml.sklearn_extensions.linear_model import StatsModelsLinearRegression @@ -30,6 +31,136 @@ ray_installed = False +@pytest.mark.parametrize("multitask", [False, True]) +@pytest.mark.parametrize("column_y", [False, True]) +@pytest.mark.parametrize("column_var", [False, True]) +@pytest.mark.parametrize("n_treatments,trimming", [(2, None), (3, None), (2, .15)]) +def test_final_sample_variance_row_scaling(multitask, column_y, column_var, n_treatments, trimming): + rng = np.random.default_rng(123) + n = 60 + X = rng.normal(size=(n, 2)) + y = rng.normal(size=n) + labels = np.arange(n) % n_treatments + T = np.eye(n_treatments)[labels, 1:] + pred = rng.normal(size=(n, n_treatments)) + propensity = rng.dirichlet(np.ones(n_treatments), size=n) + variance = rng.uniform(.5, 2, n) + weights = rng.uniform(.5, 2, n) + frequency = rng.integers(2, 5, n) + final = _ModelFinal(StatsModelsLinearRegression(), None, multitask, trimming) + final.fit(y[:, None] if column_y else y, T, X=X, + nuisances=(pred[:, None, :] if column_y else pred, propensity, propensity), + sample_weight=weights, freq_weight=frequency, + sample_var=variance[:, None] if column_var else variance) + + keep = np.ones(n, dtype=bool) if trimming is None else np.all( + (propensity >= trimming) & (propensity <= 1 - trimming), axis=1) + assert final.n_samples_trimmed_ == np.sum(~keep) + if trimming is not None: + assert 0 < np.sum(keep) < n + contrast_variance = np.array([ + [variance[i] * ((labels[i] == t) / propensity[i, t] + - (labels[i] == 0) / propensity[i, 0])**2 + for t in range(1, n_treatments)] + for i in np.flatnonzero(keep)]) + targets = pred[keep, 1:] - pred[keep, [0]][:, None] + if multitask: + expected = StatsModelsLinearRegression().fit( + X[keep], targets, sample_weight=weights[keep], freq_weight=frequency[keep], + sample_var=contrast_variance) + models = [(final.model_cate, expected)] + else: + models = [] + for t, actual in enumerate(final.models_cate): + target = targets[:, [t]] if column_y else targets[:, t] + expected = StatsModelsLinearRegression().fit( + X[keep], target, sample_weight=weights[keep], freq_weight=frequency[keep], + sample_var=contrast_variance[:, [t]] if column_y else contrast_variance[:, t]) + models.append((actual, expected)) + for actual, expected in models: + np.testing.assert_allclose(actual.predict(X), expected.predict(X)) + np.testing.assert_allclose(actual.prediction_stderr(X), expected.prediction_stderr(X)) + + +@pytest.mark.parametrize("multitask", [False, True]) +@pytest.mark.parametrize("column_y", [False, True]) +@pytest.mark.parametrize("n_treatments", [2, 3]) +def test_final_sample_variance_matches_expanded_observations(multitask, column_y, n_treatments): + rng = np.random.default_rng(456) + n = 30 + X = rng.normal(size=(n, 2)) + labels = np.arange(n) % n_treatments + T = np.eye(n_treatments)[labels, 1:] + y = X[:, 0] + labels + rng.normal(size=n) + propensity = rng.dirichlet(np.ones(n_treatments), size=n) + regression = X[:, [0]] + np.arange(n_treatments) + deviations = rng.uniform(.5, 2, (n, 1)) * [-2, 0, 2] + variance = np.var(deviations, axis=1) + weights = rng.uniform(.5, 2, n) + fits = [] + for expanded in [False, True]: + indices = np.repeat(np.arange(n), 3) if expanded else np.arange(n) + target = (y[:, None] + deviations).reshape(-1) if expanded else y + pred = regression[indices] + ( + np.eye(n_treatments)[labels[indices]] / propensity[indices] + ) * (target[:, None] - regression[indices]) + final = _ModelFinal(StatsModelsLinearRegression(), None, multitask) + final.fit(target[:, None] if column_y else target, T[indices], X=X[indices], + nuisances=(pred[:, None, :] if column_y else pred, propensity[indices], propensity[indices]), + sample_weight=weights[indices], + freq_weight=None if expanded else np.full(n, 3), + sample_var=None if expanded else (variance[:, None] if column_y else variance)) + fits.append([final.model_cate] if multitask else final.models_cate) + for summarized, expanded in zip(*fits): + np.testing.assert_allclose(summarized.predict(X), expanded.predict(X), atol=1e-10) + np.testing.assert_allclose(summarized.coef_stderr_, expanded.coef_stderr_, atol=1e-10) + np.testing.assert_allclose(summarized.intercept_stderr_, expanded.intercept_stderr_, atol=1e-10) + np.testing.assert_allclose(summarized.prediction_stderr(X), expanded.prediction_stderr(X), atol=1e-10) + + +@pytest.mark.parametrize("mode", ["linear", "separate", "multitask"]) +@pytest.mark.parametrize("n_treatments,trimming", [(3, None), (2, .25)]) +def test_drlearner_sample_variance_outcome_shapes(mode, n_treatments, trimming): + rng = np.random.default_rng(123) + n = 150 + X = rng.normal(size=(n, 2)) + T = rng.permutation(np.arange(n) % n_treatments) + y = T + X[:, 0] + rng.normal(size=n) + variance = rng.uniform(.5, 2, n) + weights = rng.uniform(.5, 2, n) + frequency = rng.integers(2, 5, n) + results = [] + for column_y, column_var in [(False, False), (True, False), (True, True), (False, True)]: + kwargs = dict(model_regression=LinearRegression(), model_propensity=LogisticRegression(), + random_state=123, trimming_threshold=trimming) + est = LinearDRLearner(**kwargs) if mode == "linear" else DRLearner( + **kwargs, model_final=StatsModelsLinearRegression(), multitask_model_final=mode == "multitask") + est.fit(y[:, None] if column_y else y, T, X=X, sample_weight=weights, + freq_weight=frequency, sample_var=variance[:, None] if column_var else variance) + effects = est.const_marginal_effect(X) + assert effects.shape == ((n, 1, n_treatments - 1) if column_y else (n, n_treatments - 1)) + if mode == "multitask": + stderr = est.multitask_model_cate.prediction_stderr(X) + else: + stderr = est.const_marginal_effect_inference(X).stderr + results.append((effects.reshape(n, -1), stderr.reshape(n, -1))) + for effects, stderr in results[1:]: + np.testing.assert_allclose(effects, results[0][0], atol=1e-10) + np.testing.assert_allclose(stderr, results[0][1], atol=1e-10) + + +@pytest.mark.parametrize("multitask", [False, True]) +@pytest.mark.parametrize("invalid_shape", [(12, 2), (12, 1, 1), (1, 12), (11,)]) +def test_final_rejects_invalid_sample_variance_shape(multitask, invalid_shape): + final = _ModelFinal(StatsModelsLinearRegression(), None, multitask) + y = np.arange(12, dtype=float) + T = (y % 2)[:, None] + with pytest.raises(ValueError, match="sample_var must have shape"): + final.fit(y, T, X=y[:, None], + nuisances=(np.column_stack([y, y + 1]), np.full((12, 2), .5), np.full((12, 2), .5)), + freq_weight=np.full(12, 2), sample_var=np.ones(invalid_shape)) + + @pytest.mark.serial class TestDRLearner(unittest.TestCase): diff --git a/econml/tests/test_inference.py b/econml/tests/test_inference.py index 4d19d2fd1..2d7f211ad 100644 --- a/econml/tests/test_inference.py +++ b/econml/tests/test_inference.py @@ -6,6 +6,9 @@ import unittest import pytest import pickle +from unittest.mock import patch +from scipy.optimize import brentq +from scipy.stats import norm from sklearn.base import clone from sklearn.preprocessing import PolynomialFeatures from sklearn.linear_model import LinearRegression, LogisticRegression, Lasso @@ -17,6 +20,74 @@ from econml.utilities import get_feature_names_or_default, get_input_columns +@pytest.mark.parametrize("output_shape", [(), (1,), (2,), (1, 3), (2, 3)]) +@pytest.mark.parametrize("zero_scales", ["none", "sample", "column", "all"]) +def test_population_mixture_zero_scales(output_shape, zero_scales): + n_columns = int(np.prod(output_shape)) + mean = np.arange(3 * n_columns, dtype=float).reshape(3, n_columns) + 10 + stderr = np.linspace(.5, 2, 3 * n_columns).reshape(3, n_columns) + if zero_scales == "sample": + stderr[1, 0] = 0 + elif zero_scales == "column": + stderr[:, 0] = 0 + elif zero_scales == "all": + stderr[:] = 0 + + expected = np.full((2, n_columns), np.nan) + for j in range(n_columns): + if np.any(stderr[:, j] == 0): + continue + for i, alpha in enumerate([.025, .975]): + bounds = norm.ppf(alpha, loc=mean[:, j], scale=stderr[:, j]) + expected[i, j] = brentq( + lambda x: np.mean(norm.cdf(x, loc=mean[:, j], scale=stderr[:, j])) - alpha, + np.min(bounds), np.max(bounds)) + + for order in [np.arange(n_columns), np.arange(n_columns)[::-1]]: + pop = PopulationSummaryResults( + mean[:, order].reshape((3,) + output_shape), + stderr[:, order].reshape((3,) + output_shape), + None, d_t=output_shape[-1] if len(output_shape) == 2 else 1, + d_y=output_shape[0] if output_shape else 1) + intervals = pop.conf_int_point(tol=1e-9) + for i, interval in enumerate(intervals): + assert interval.shape == (1,) + output_shape + np.testing.assert_allclose( + interval.reshape(-1), expected[i, order], atol=1e-6, equal_nan=True) + + +@pytest.mark.parametrize("output_shape", [(3,), (1, 3)]) +@pytest.mark.parametrize("field,value", [ + ("mean", np.nan), ("mean", np.inf), ("mean", -np.inf), + ("stderr", np.nan), ("stderr", np.inf), ("stderr", -np.inf), ("stderr", -1), +]) +def test_population_mixture_invalid_columns(output_shape, field, value): + mean = np.array([[0., 1., 10.], [0., 2., 12.], [0., 3., 14.]]) + stderr = np.ones_like(mean) + stderr[:, 0] = 0 + (mean if field == "mean" else stderr)[1, 1] = value + reference = PopulationSummaryResults(mean[:, 2], stderr[:, 2], None, d_t=1, d_y=1) + expected = reference.conf_int_point(tol=1e-9) + original_ppf = norm.ppf + + def checked_ppf(alpha, *, loc, scale): + # Fail promptly rather than hanging if invalid parameters reach the iterative solver. + assert np.all(np.isfinite(loc)) + assert np.all(np.isfinite(scale) & (scale > 0)) + return original_ppf(alpha, loc=loc, scale=scale) + + for order in [np.arange(3), np.array([2, 0, 1])]: + pop = PopulationSummaryResults( + mean[:, order].reshape((3,) + output_shape), + stderr[:, order].reshape((3,) + output_shape), None, + d_t=output_shape[-1] if len(output_shape) == 2 else 1, d_y=output_shape[0]) + with patch("econml.inference._inference.scipy.stats.norm.ppf", side_effect=checked_ppf): + intervals = pop.conf_int_point(tol=1e-9) + for interval, bound in zip(intervals, expected): + assert interval.shape == (1,) + output_shape + np.testing.assert_allclose(interval.reshape(-1), np.array([np.nan, np.nan, bound.item()])[order]) + + class TestInference(unittest.TestCase): @classmethod diff --git a/econml/tests/test_model_selection.py b/econml/tests/test_model_selection.py index a08373a40..c70f9d1c7 100644 --- a/econml/tests/test_model_selection.py +++ b/econml/tests/test_model_selection.py @@ -3,19 +3,66 @@ import unittest import numpy as np +import pytest from scipy.special import expit from sklearn.ensemble import RandomForestRegressor from sklearn.linear_model import (ElasticNetCV, Lasso, LassoCV, LinearRegression, LogisticRegression, LogisticRegressionCV, MultiTaskElasticNetCV, MultiTaskLassoCV, RidgeCV, RidgeClassifierCV) -from sklearn.model_selection import GridSearchCV, RandomizedSearchCV +from sklearn.model_selection import GridSearchCV, KFold, RandomizedSearchCV from sklearn.preprocessing import PolynomialFeatures from econml.dml import LinearDML from econml.sklearn_extensions.linear_model import WeightedLassoCVWrapper +from econml.sklearn_extensions.model_selection import get_selector from econml.utilities import MultiModelWrapper from econml.dr import LinearDRLearner +@pytest.mark.parametrize("model,weighted", [ + (MultiTaskElasticNetCV(), False), + (MultiTaskLassoCV(), False), + (WeightedLassoCVWrapper(), False), + (WeightedLassoCVWrapper(), True), +]) +def test_multitask_selector_score_output_variance(model, weighted): + rng = np.random.default_rng(123) + X = rng.normal(size=(150, 3)) + y = rng.normal(size=(150, 2)) * [1, 3] + [10, -20] + folds = list(KFold(3).split(X)) + kwargs = {"sample_weight": rng.uniform(.5, 2, len(y))} if weighted else {} + scores = [] + mse_paths = [] + for target in [y, y + [1000, -300], y[:, ::-1], np.tile(y, (1, 2))]: + selector = get_selector(model, is_discrete=False) + selector.train(True, folds, X, target, **kwargs) + min_mse = np.min(np.mean(selector.searcher.mse_path_, axis=-1)) + expected_variance = np.mean([np.var(target[:, j]) for j in range(target.shape[1])]) + np.testing.assert_allclose(selector.best_score, 1 - min_mse / expected_variance) + scores.append(selector.best_score) + mse_paths.append(selector.searcher.mse_path_) + np.testing.assert_allclose(scores[:3], scores[0], atol=1e-10) + for path in mse_paths[1:3]: + np.testing.assert_allclose(path, mse_paths[0], atol=1e-10) + + +@pytest.mark.parametrize("model", [ElasticNetCV(), LassoCV(), WeightedLassoCVWrapper()]) +@pytest.mark.parametrize("weighted", [False, True]) +def test_single_output_selector_score_variance(model, weighted): + rng = np.random.default_rng(123) + X = rng.normal(size=(150, 3)) + y = X[:, 0] + rng.normal(size=150) + folds = list(KFold(3).split(X)) + kwargs = {"sample_weight": rng.uniform(.5, 2, len(y))} if weighted else {} + scores = [] + for target in [y, y[:, None], y + 1000]: + selector = get_selector(model, is_discrete=False) + selector.train(True, folds, X, target, **kwargs) + min_mse = np.min(np.mean(selector.searcher.mse_path_, axis=-1)) + np.testing.assert_allclose(selector.best_score, 1 - min_mse / np.var(target)) + scores.append(selector.best_score) + np.testing.assert_allclose(scores, scores[0], atol=1e-10) + + class TestModelSelection(unittest.TestCase): def _simple_dgp(self, n, d_x, d_w, discrete_treatment):