"Fixes linear DRIV shape inconsistency"

Signed-off-by: ssemov <ssemov@gmail.com>

econml/inference/_inference.py

@@ -13,7 +13,8 @@
 from ._bootstrap import BootstrapEstimator
 from ..utilities import (Summary, _safe_norm_ppf, broadcast_unit_treatments,
                          cross_product, inverse_onehot, ndim,
-                         parse_final_model_params, reshape_treatmentwise_effects, shape, filter_none_kwargs)
+                         parse_final_model_params, reshape_treatmentwise_effects, shape, filter_none_kwargs,
+                         reshape_outcomewise_effects)
 
 """Options for performing inference in estimators."""
 
@@ -481,6 +482,29 @@ def prefit(self, estimator, *args, **kwargs):
                                                       "their features, but here fit_intercept is True")
         self.model_final.cov_type = self.cov_type
 
+    def effect_inference(self, X, *, T0, T1):
+        # We can write effect inference as a function of prediction and prediction standard error of
+        # the final method for linear models
+        X, T0, T1 = self._est._expand_treatments(X, T0, T1)
+        if X is None:
+            X = np.ones((T0.shape[0], 1))
+        elif self.featurizer is not None:
+            X = self.featurizer.transform(X)
+        XT = cross_product(X, T1 - T0)
+        e_pred = reshape_outcomewise_effects(self._predict(XT), self._d_y)
+        e_stderr = reshape_outcomewise_effects(self._prediction_stderr(XT), self._d_y)
+        d_y = self._d_y[0] if self._d_y else 1
+
+        mean_XT = XT.mean(axis=0, keepdims=True)
+        mean_pred_stderr = self._prediction_stderr(mean_XT)  # shape[0] will always be 1 here
+        # squeeze the first axis
+        mean_pred_stderr = np.squeeze(mean_pred_stderr, axis=0) if mean_pred_stderr is not None else None
+        # d_t=None here since we measure the effect across all Ts
+        return NormalInferenceResults(d_t=None, d_y=d_y, pred=e_pred,
+                                      pred_stderr=e_stderr, mean_pred_stderr=mean_pred_stderr, inf_type='effect',
+                                      feature_names=self._est.cate_feature_names(),
+                                      output_names=self._est.cate_output_names())
+
 
 class GenericModelFinalInferenceDiscrete(Inference):
     """

econml/utilities.py

@@ -770,6 +770,28 @@ def reshape_treatmentwise_effects(A, d_t, d_y):
     else:
         return A
 
+def reshape_outcomewise_effects(A, d_y):
+    """
+    Given an effects matrix, reshape second dimension to be consistent with d_y[0].
+
+    Parameters
+    ----------
+    A : array
+        The effects array to be reshaped. It should have shape (m,) or (m, d_y).
+    d_y : tuple of int
+        Either () if Y was a vector, or a 1-tuple of the number of columns of Y if it was an array.
+
+    Returns
+    -------
+    A : array
+        The reshaped effects array with shape:
+         - (m, ) if d_y is () and Y is a vector,
+         - (m, d_y) if d_y is a 1-tuple and Y is an array.
+    """
+    if np.shape(A)[1:] == d_y or d_y == ():
+        return A
+    else:
+        return A.reshape(-1, d_y[0])
 
 def einsum_sparse(subscripts, *arrs):
     """