Fix output-axis reductions and DR variance propagation

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
@@ -1328,7 +1337,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
@@ -1342,9 +1351,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

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)

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
 
 

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):