Merge pull request #7 from PierreBoyeau/ppi_mean_multid

[WIP] fix multid bug + fast `_calc_lhat_glm`

ppi_py/ppi.py

@@ -7,6 +7,7 @@
 from statsmodels.stats.weightstats import _zconfint_generic, _zstat_generic
 from sklearn.linear_model import LogisticRegression
 from .utils import (
+ construct_weight_vector,
  safe_expit,
  safe_log1pexp,
  compute_cdf,
@@ -15,6 +16,7 @@
  linfty_dkw,
  linfty_binom,
  form_discrete_distribution,
+ reshape_to_2d,
 )
 
 
@@ -55,7 +57,14 @@ def rectified_p_value(
 
 
 def ppi_mean_pointestimate(
- Y, Yhat, Yhat_unlabeled, lhat=None, coord=None, w=None, w_unlabeled=None
+ Y,
+ Yhat,
+ Yhat_unlabeled,
+ lhat=None,
+ coord=None,
+ w=None,
+ w_unlabeled=None,
+ lambd_optim_mode="overall",
 ):
  """Computes the prediction-powered point estimate of the d-dimensional mean.
 
@@ -74,19 +83,20 @@ def ppi_mean_pointestimate(
  Notes:
  `[ADZ23] <https://arxiv.org/abs/2311.01453>`__ A. N. Angelopoulos, J. C. Duchi, and T. Zrnic. PPI++: Efficient Prediction Powered Inference. arxiv:2311.01453, 2023.
  """
+ Y = reshape_to_2d(Y)
+ Yhat = reshape_to_2d(Yhat)
+ Yhat_unlabeled = reshape_to_2d(Yhat_unlabeled)
  n = Y.shape[0]
  N = Yhat_unlabeled.shape[0]
- d = Yhat.shape[1] if len(Yhat.shape) > 1 else 1
- w = np.ones(n) if w is None else w / w.sum() * n
- w_unlabeled = (
- np.ones(N)
- if w_unlabeled is None
- else w_unlabeled / w_unlabeled.sum() * N
- )
+ d = Yhat.shape[1]
+
+ w = construct_weight_vector(n, w, vectorized=True)
+ w_unlabeled = construct_weight_vector(N, w_unlabeled, vectorized=True)
+
  if lhat is None:
- ppi_pointest = (w_unlabeled * Yhat_unlabeled).mean() + (
+ ppi_pointest = (w_unlabeled * Yhat_unlabeled).mean(0) + (
  w * (Y - Yhat)
- ).mean()
+ ).mean(0)
  grads = w * (Y - ppi_pointest)
  grads_hat = w * (Yhat - ppi_pointest)
  grads_hat_unlabeled = w_unlabeled * (Yhat_unlabeled - ppi_pointest)
@@ -98,6 +108,7 @@ def ppi_mean_pointestimate(
  inv_hessian,
  coord=None,
  clip=True,
+ optim_mode=lambd_optim_mode,
  )
  return ppi_mean_pointestimate(
  Y,
@@ -111,7 +122,7 @@ def ppi_mean_pointestimate(
  else:
  return (w_unlabeled * lhat * Yhat_unlabeled).mean(axis=0) + (
  w * (Y - lhat * Yhat)
- ).mean(axis=0)
+ ).mean(axis=0).squeeze()
 
 
 def ppi_mean_ci(
@@ -124,6 +135,7 @@ def ppi_mean_ci(
  coord=None,
  w=None,
  w_unlabeled=None,
+ lambd_optim_mode="overall",
 ):
  """Computes the prediction-powered confidence interval for a d-dimensional mean.
 
@@ -147,12 +159,13 @@ def ppi_mean_ci(
  n = Y.shape[0]
  N = Yhat_unlabeled.shape[0]
  d = Y.shape[1] if len(Y.shape) > 1 else 1
- w = np.ones(n) if w is None else w / w.sum() * n
- w_unlabeled = (
- np.ones(N)
- if w_unlabeled is None
- else w_unlabeled / w_unlabeled.sum() * N
- )
+
+ Y = reshape_to_2d(Y)
+ Yhat = reshape_to_2d(Yhat)
+ Yhat_unlabeled = reshape_to_2d(Yhat_unlabeled)
+
+ w = construct_weight_vector(n, w, vectorized=True)
+ w_unlabeled = construct_weight_vector(N, w_unlabeled, vectorized=True)
 
  if lhat is None:
  ppi_pointest = ppi_mean_pointestimate(
@@ -174,6 +187,7 @@ def ppi_mean_ci(
  inv_hessian,
  coord=None,
  clip=True,
+ optim_mode=lambd_optim_mode,
  )
  return ppi_mean_ci(
  Y,
@@ -196,8 +210,8 @@ def ppi_mean_ci(
  w_unlabeled=w_unlabeled,
  )
 
- imputed_std = (w_unlabeled * (lhat * Yhat_unlabeled)).std() / np.sqrt(N)
- rectifier_std = (w * (Y - lhat * Yhat)).std() / np.sqrt(n)
+ imputed_std = (w_unlabeled * (lhat * Yhat_unlabeled)).std(0) / np.sqrt(N)
+ rectifier_std = (w * (Y - lhat * Yhat)).std(0) / np.sqrt(n)
 
  return _zconfint_generic(
  ppi_pointest,
@@ -217,6 +231,7 @@ def ppi_mean_pval(
  coord=None,
  w=None,
  w_unlabeled=None,
+ lambd_optim_mode="overall",
 ):
  """Computes the prediction-powered p-value for a 1D mean.
 
@@ -239,35 +254,39 @@ def ppi_mean_pval(
  """
  n = Y.shape[0]
  N = Yhat_unlabeled.shape[0]
- w = np.ones(n) if w is None else w / w.sum() * n
- w_unlabeled = (
- np.ones(N)
- if w_unlabeled is None
- else w_unlabeled / w_unlabeled.sum() * N
- )
+ # w = np.ones(n) if w is None else w / w.sum() * n
+ w = construct_weight_vector(n, w, vectorized=True)
+ w_unlabeled = construct_weight_vector(N, w_unlabeled, vectorized=True)
+
+ Y = reshape_to_2d(Y)
+ Yhat = reshape_to_2d(Yhat)
+ Yhat_unlabeled = reshape_to_2d(Yhat_unlabeled)
+ d = Y.shape[1]
 
  if lhat is None:
- if len(Y.shape) > 1 and Y.shape[1] > 1:
- lhat = 1
- else:
- ppi_pointest = (w_unlabeled * Yhat_unlabeled).mean() + (
- w * (Y - Yhat)
- ).mean()
- grads = w * (Y - ppi_pointest)
- grads_hat = w * (Yhat - ppi_pointest)
- grads_hat_unlabeled = w_unlabeled * (Yhat_unlabeled - ppi_pointest)
- inv_hessian = np.ones((1, 1))
- lhat = _calc_lhat_glm(
- grads, grads_hat, grads_hat_unlabeled, inv_hessian, coord=None
- )
+ ppi_pointest = (w_unlabeled * Yhat_unlabeled).mean(0) + (
+ w * (Y - Yhat)
+ ).mean(0)
+ grads = w * (Y - ppi_pointest)
+ grads_hat = w * (Yhat - ppi_pointest)
+ grads_hat_unlabeled = w_unlabeled * (Yhat_unlabeled - ppi_pointest)
+ inv_hessian = np.eye(d)
+ lhat = _calc_lhat_glm(
+ grads,
+ grads_hat,
+ grads_hat_unlabeled,
+ inv_hessian,
+ coord=None,
+ optim_mode=lambd_optim_mode,
+ )
 
  return rectified_p_value(
- (w * Y - lhat * w * Yhat).mean(),
- (w * Y - lhat * w * Yhat).std() / np.sqrt(n),
- (w_unlabeled * lhat * Yhat_unlabeled).mean(),
- (w_unlabeled * lhat * Yhat_unlabeled).std() / np.sqrt(N),
- null,
- alternative,
+ rectifier=(w * Y - lhat * w * Yhat).mean(0),
+ rectifier_std=(w * Y - lhat * w * Yhat).std(0) / np.sqrt(n),
+ imputed_mean=(w_unlabeled * lhat * Yhat_unlabeled).mean(0),
+ imputed_std=(w_unlabeled * lhat * Yhat_unlabeled).std(0) / np.sqrt(N),
+ null=null,
+ alternative=alternative,
  )
 
 
@@ -1049,7 +1068,13 @@ def ppi_logistic_ci(
 
 
 def _calc_lhat_glm(
- grads, grads_hat, grads_hat_unlabeled, inv_hessian, coord=None, clip=False
+ grads,
+ grads_hat,
+ grads_hat_unlabeled,
+ inv_hessian,
+ coord=None,
+ clip=False,
+ optim_mode="overall",
 ):
  """
  Calculates the optimal value of lhat for the prediction-powered confidence interval for GLMs.
@@ -1059,38 +1084,41 @@ def _calc_lhat_glm(
  grads_hat (ndarray): Gradient of the loss function with respect to the model parameter evaluated using predictions on the labeled data.
  grads_hat_unlabeled (ndarray): Gradient of the loss function with respect to the parameter evaluated using predictions on the unlabeled data.
  inv_hessian (ndarray): Inverse of the Hessian of the loss function with respect to the parameter.
- coord (int, optional): Coordinate for which to optimize `lhat`. If `None`, it optimizes the total variance over all coordinates. Must be in {1, ..., d} where d is the shape of the estimand.
+ coord (int, optional): Coordinate for which to optimize `lhat`, when `optim_mode="overall"`.
+ If `None`, it optimizes the total variance over all coordinates. Must be in {1, ..., d} where d is the shape of the estimand.
  clip (bool, optional): Whether to clip the value of lhat to be non-negative. Defaults to `False`.
+ optim_mode (ndarray, optional): Mode for which to optimize `lhat`, either `overall` or `element`.
+ If `overall`, it optimizes the total variance over all coordinates, and the function returns a scalar.
+ If `element`, it optimizes the variance for each coordinate separately, and the function returns a vector.
+
 
  Returns:
  float: Optimal value of `lhat`. Lies in [0,1].
  """
+ grads = reshape_to_2d(grads)
+ grads_hat = reshape_to_2d(grads_hat)
+ grads_hat_unlabeled = reshape_to_2d(grads_hat_unlabeled)
  n = grads.shape[0]
  N = grads_hat_unlabeled.shape[0]
  d = inv_hessian.shape[0]
- cov_grads = np.zeros((d, d))
-
- for i in range(n):
- cov_grads += (1 / n) * (
- np.outer(
- grads[i] - grads.mean(axis=0),
- grads_hat[i] - grads_hat.mean(axis=0),
- )
- + np.outer(
- grads_hat[i] - grads_hat.mean(axis=0),
- grads[i] - grads.mean(axis=0),
- )
+ if grads.shape[1] != d:
+ raise ValueError(
+ "Dimension mismatch between the gradient and the inverse Hessian."
  )
+
+ grads_cent = grads - grads.mean(axis=0)
+ grad_hat_cent = grads_hat - grads_hat.mean(axis=0)
+ cov_grads = (1 / n) * (
+ grads_cent.T @ grad_hat_cent + grad_hat_cent.T @ grads_cent
+ )
+
  var_grads_hat = np.cov(
  np.concatenate([grads_hat, grads_hat_unlabeled], axis=0).T
  )
+ var_grads_hat = var_grads_hat.reshape(d, d)
 
- if coord is None:
- vhat = inv_hessian
- else:
- vhat = inv_hessian @ np.eye(d)[coord]
-
- if d > 1:
+ vhat = inv_hessian if coord is None else inv_hessian[coord, coord]
+ if optim_mode == "overall":
  num = (
  np.trace(vhat @ cov_grads @ vhat)
  if coord is None
@@ -1101,14 +1129,19 @@ def _calc_lhat_glm(
  if coord is None
  else 2 * (1 + (n / N)) * vhat @ var_grads_hat @ vhat
  )
+ lhat = num / denom
+ lhat = lhat.item()
+ elif optim_mode == "element":
+ num = np.diag(vhat @ cov_grads @ vhat)
+ denom = 2 * (1 + (n / N)) * np.diag(vhat @ var_grads_hat @ vhat)
+ lhat = num / denom
  else:
- num = vhat * cov_grads * vhat
- denom = 2 * (1 + (n / N)) * vhat * var_grads_hat * vhat
-
- lhat = num / denom
+ raise ValueError(
+ "Invalid value for optim_mode. Must be either 'overall' or 'element'."
+ )
  if clip:
  lhat = np.clip(lhat, 0, 1)
- return lhat.item()
+ return lhat
 
 
 """

ppi_py/utils/statistics_utils.py

@@ -4,6 +4,22 @@
 from scipy.optimize import brentq
 
 
+def construct_weight_vector(n_obs, existing_weight, vectorized=False):
+ res = (
+ np.ones(n_obs)
+ if existing_weight is None
+ else existing_weight / existing_weight.sum() * n_obs
+ )
+ if vectorized and (len(res.shape) == 1):
+ res = res[:, None]
+ return res
+
+
+def reshape_to_2d(x):
+ """Reshapes a 1D array to a 2D array."""
+ return x.reshape(-1, 1) if len(x.shape) == 1 else x.copy()
+
+
 @njit
 def safe_expit(x):
  """Computes the sigmoid function in a numerically stable way."""

tests/test_mean.py

@@ -30,6 +30,49 @@ def test_ppi_mean_ci():
  assert not failed
 
 
+def test_ppi_mean_multid():
+ trials = 100
+ alphas = np.array([0.5, 0.2, 0.1, 0.05, 0.01])
+ n_alphas = alphas.shape[0]
+ n_dims = 5
+ epsilon = 0.1
+ includeds = np.zeros((n_alphas, n_dims))
+ for _ in range(trials):
+ Y = np.random.normal(0, 1, (10000, n_dims))
+ Yhat = np.random.normal(-2, 1, (10000, n_dims))
+ Yhat_unlabeled = np.random.normal(-2, 1, (10000, n_dims))
+ for j in range(alphas.shape[0]):
+ ci = ppi_mean_ci(Y, Yhat, Yhat_unlabeled, alpha=alphas[j])
+
+ included = (ci[0] <= 0) & (ci[1] >= 0)
+ includeds[j] += included.astype(int)
+ failed = np.any(includeds / trials < 1 - alphas - epsilon)
+ assert not failed
+
+
+def test_ppi_mean_elem():
+ alpha = 0.1
+ Y = np.random.normal(0, 1, 10000)
+ Yhat = np.random.normal(-2, 1, 10000)
+ Yhat_unlabeled = np.random.normal(-2, 1, 10000)
+
+ ppi_mean_pointestimate(Y, Yhat, Yhat_unlabeled, lambd_optim_mode="element")
+ ppi_mean_ci(
+ Y, Yhat, Yhat_unlabeled, alpha=alpha, lambd_optim_mode="element"
+ )
+ ppi_mean_pval(Y, Yhat, Yhat_unlabeled, lambd_optim_mode="element")
+
+ Y = np.random.normal(0, 1, (10000, 5))
+ Yhat = np.random.normal(-2, 1, (10000, 5))
+ Yhat_unlabeled = np.random.normal(-2, 1, (10000, 5))
+
+ ppi_mean_pointestimate(Y, Yhat, Yhat_unlabeled, lambd_optim_mode="element")
+ ppi_mean_ci(
+ Y, Yhat, Yhat_unlabeled, alpha=alpha, lambd_optim_mode="element"
+ )
+ ppi_mean_pval(Y, Yhat, Yhat_unlabeled, lambd_optim_mode="element")
+
+
 def test_ppi_mean_pval():
  trials = 1000
  alphas = np.array([0.5, 0.2, 0.1, 0.05, 0.01])