Optimizing NormalInferenceResults confidence interval method speed (#879)

* Fixed normal inference results confidence interval unnecessary loop

Signed-off-by: gdaiha <[email protected]>
Signed-off-by: Gabriel Daiha <[email protected]>

Author: gdaiha

econml/inference/_inference.py

@@ -1066,14 +1066,9 @@ def conf_int(self, alpha=0.05):
         """
         if self.stderr is None:
             raise AttributeError("Only point estimates are available!")
-        if np.isscalar(self.point_estimate):
+        else:
             return _safe_norm_ppf(alpha / 2, loc=self.point_estimate, scale=self.stderr), \
                 _safe_norm_ppf(1 - alpha / 2, loc=self.point_estimate, scale=self.stderr)
-        else:
-            return np.array([_safe_norm_ppf(alpha / 2, loc=p, scale=err)
-                             for p, err in zip(self.point_estimate, self.stderr)]), \
-                np.array([_safe_norm_ppf(1 - alpha / 2, loc=p, scale=err)
-                          for p, err in zip(self.point_estimate, self.stderr)])
 
     def pvalue(self, value=0):
         """
@@ -1398,14 +1393,8 @@ def conf_int_mean(self, *, alpha=None):
         alpha = self.alpha if alpha is None else alpha
         mean_point = self.mean_point
         stderr_mean = self.stderr_mean
-        if np.isscalar(mean_point):
-            return (_safe_norm_ppf(alpha / 2, loc=mean_point, scale=stderr_mean),
-                    _safe_norm_ppf(1 - alpha / 2, loc=mean_point, scale=stderr_mean))
-        else:
-            return np.array([_safe_norm_ppf(alpha / 2, loc=p, scale=err)
-                             for p, err in zip(mean_point, stderr_mean)]), \
-                np.array([_safe_norm_ppf(1 - alpha / 2, loc=p, scale=err)
-                          for p, err in zip(mean_point, stderr_mean)])
+        return (_safe_norm_ppf(alpha / 2, loc=mean_point, scale=stderr_mean),
+                _safe_norm_ppf(1 - alpha / 2, loc=mean_point, scale=stderr_mean))
 
     @property
     def std_point(self):

econml/sklearn_extensions/linear_model.py

@@ -1627,10 +1627,8 @@ def coef__interval(self, alpha=0.05):
         coef__interval : {tuple ((p, d) array, (p,d) array), tuple ((d,) array, (d,) array)}
             The lower and upper bounds of the confidence interval of the coefficients
         """
-        return np.array([_safe_norm_ppf(alpha / 2, loc=p, scale=err)
-                         for p, err in zip(self.coef_, self.coef_stderr_)]), \
-            np.array([_safe_norm_ppf(1 - alpha / 2, loc=p, scale=err)
-                      for p, err in zip(self.coef_, self.coef_stderr_)])
+        return (_safe_norm_ppf(alpha / 2, loc=self.coef_, scale=self.coef_stderr_),
+                _safe_norm_ppf(1 - alpha / 2, loc=self.coef_, scale=self.coef_stderr_))
 
     def intercept__interval(self, alpha=0.05):
         """
@@ -1651,14 +1649,8 @@ def intercept__interval(self, alpha=0.05):
             return (0 if self._n_out == 0 else np.zeros(self._n_out)), \
                 (0 if self._n_out == 0 else np.zeros(self._n_out))
 
-        if self._n_out == 0:
-            return _safe_norm_ppf(alpha / 2, loc=self.intercept_, scale=self.intercept_stderr_), \
-                _safe_norm_ppf(1 - alpha / 2, loc=self.intercept_, scale=self.intercept_stderr_)
-        else:
-            return np.array([_safe_norm_ppf(alpha / 2, loc=p, scale=err)
-                             for p, err in zip(self.intercept_, self.intercept_stderr_)]), \
-                np.array([_safe_norm_ppf(1 - alpha / 2, loc=p, scale=err)
-                          for p, err in zip(self.intercept_, self.intercept_stderr_)])
+        return (_safe_norm_ppf(alpha / 2, loc=self.intercept_, scale=self.intercept_stderr_),
+                _safe_norm_ppf(1 - alpha / 2, loc=self.intercept_, scale=self.intercept_stderr_))
 
     def predict_interval(self, X, alpha=0.05):
         """
@@ -1677,10 +1669,12 @@ def predict_interval(self, X, alpha=0.05):
         prediction_intervals : {tuple ((n,) array, (n,) array), tuple ((n,p) array, (n,p) array)}
             The lower and upper bounds of the confidence intervals of the predicted mean outcomes
         """
-        return np.array([_safe_norm_ppf(alpha / 2, loc=p, scale=err)
-                         for p, err in zip(self.predict(X), self.prediction_stderr(X))]), \
-            np.array([_safe_norm_ppf(1 - alpha / 2, loc=p, scale=err)
-                      for p, err in zip(self.predict(X), self.prediction_stderr(X))])
+
+        pred = self.predict(X)
+        pred_stderr = self.prediction_stderr(X)
+
+        return (_safe_norm_ppf(alpha / 2, loc=pred, scale=pred_stderr),
+                _safe_norm_ppf(1 - alpha / 2, loc=pred, scale=pred_stderr))
 
 
 class StatsModelsLinearRegression(_StatsModelsWrapper):