cbps 0.2.0__py3-none-any.whl

cbps/diagnostics/weights_diag.py

@@ -0,0 +1,205 @@
+"""
+Weight Quality Diagnostics
+===========================
+
+Comprehensive diagnostics for inverse probability weights produced by CBPS
+estimation, including effective sample size (ESS), weight distribution
+summaries, and extreme value detection.
+
+The Kish (1965) effective sample size is the primary metric for assessing
+whether extreme weights are degrading estimation precision.
+
+References
+----------
+Kish, L. (1965). Survey Sampling. Wiley, New York.
+
+Imai, K. and Ratkovic, M. (2014). Covariate balancing propensity score.
+Journal of the Royal Statistical Society, Series B, 76(1), 243-263.
+"""
+
+import numpy as np
+from typing import Optional
+
+
+def weight_diagnostics(weights, treat=None):
+    """Compute comprehensive weight quality diagnostics.
+
+    Based on Kish (1965) effective sample size and standard
+    IPW weight quality indicators.
+
+    Parameters
+    ----------
+    weights : np.ndarray
+        IPW weights from CBPS estimation. Should be non-negative for
+        meaningful ESS interpretation. If negative weights are present
+        (e.g., from ATT balance conditions), a warning is issued and
+        ESS is computed on absolute values.
+    treat : np.ndarray, optional
+        Treatment indicator for group-specific diagnostics.
+
+    Returns
+    -------
+    dict with:
+        - ess: Kish effective sample size
+        - ess_ratio: ESS / n (closer to 1 = better)
+        - weight_max: maximum absolute weight
+        - weight_min: minimum absolute weight (among nonzero)
+        - weight_ratio: max/min ratio (of absolute values)
+        - cv: coefficient of variation of weights
+        - n_extreme: count of weights with abs(w) > 10*median(abs(w))
+        - n_negative: count of negative weights (0 if all non-negative)
+        - warning_level: 'ok'/'caution'/'severe'
+        - group_diagnostics: dict per treatment group (if treat provided)
+
+    Notes
+    -----
+    Warning thresholds (based on Kish 1965, Chapter 11):
+        - ESS/n < 0.5 → 'caution'
+        - ESS/n < 0.2 → 'severe'
+
+    The ESS formula is: ESS = (sum(w))^2 / sum(w^2)
+    For uniform weights, ESS = n. For highly variable weights, ESS << n.
+
+    When negative weights are present, the Kish ESS formula does not have
+    its standard interpretation. In this case, ESS is computed on abs(w) and
+    a warning is included in the result.
+
+    References
+    ----------
+    Kish, L. (1965). Survey Sampling. Wiley, New York. Chapter 11.
+    """
+    weights = np.asarray(weights, dtype=float).ravel()
+    n = len(weights)
+
+    # Handle degenerate cases
+    if n == 0:
+        return {
+            'ess': 0.0,
+            'ess_ratio': 0.0,
+            'weight_max': np.nan,
+            'weight_min': np.nan,
+            'weight_ratio': np.nan,
+            'cv': np.nan,
+            'n_extreme': 0,
+            'n_negative': 0,
+            'warning_level': 'severe',
+            'group_diagnostics': None,
+        }
+
+    # Detect negative weights
+    n_negative = int(np.sum(weights < 0))
+    has_negative = n_negative > 0
+
+    # For ESS computation: use absolute values when negative weights present
+    # Kish ESS is only interpretable for non-negative weights
+    if has_negative:
+        import warnings
+        warnings.warn(
+            f"Kish ESS is defined for non-negative weights. "
+            f"{n_negative} negative weight(s) detected; "
+            f"ESS is computed on |weights| as an approximation. "
+            f"Consider using only the final IPW weights (not balance weights) "
+            f"for this diagnostic.",
+            UserWarning,
+            stacklevel=2
+        )
+        w_for_ess = np.abs(weights)
+    else:
+        w_for_ess = weights
+
+    sum_w = np.sum(w_for_ess)
+    sum_w2 = np.sum(w_for_ess ** 2)
+
+    # ESS computation (Kish 1965)
+    if sum_w2 == 0:
+        # All weights are zero
+        ess = 0.0
+        ess_ratio = 0.0
+    else:
+        ess = (sum_w ** 2) / sum_w2
+        ess_ratio = ess / n
+
+    # Weight range based on absolute values (captures extreme negative weights)
+    abs_weights = np.abs(weights)
+    nonzero_mask = abs_weights > 0
+    if np.any(nonzero_mask):
+        weight_min = float(np.min(abs_weights[nonzero_mask]))
+        weight_max = float(np.max(abs_weights[nonzero_mask]))
+    else:
+        weight_min = 0.0
+        weight_max = 0.0
+
+    # Max/min ratio
+    if weight_min > 0:
+        weight_ratio = weight_max / weight_min
+    else:
+        weight_ratio = np.inf if weight_max > 0 else np.nan
+
+    # Coefficient of variation (on absolute values when negative present)
+    w_for_cv = w_for_ess
+    w_mean = np.mean(w_for_cv)
+    if w_mean > 0:
+        cv = float(np.std(w_for_cv) / w_mean)
+    else:
+        cv = np.nan
+
+    # Extreme weight count: abs(w) > 10 * median(abs(w))
+    median_abs_w = np.median(abs_weights)
+    if median_abs_w > 0:
+        n_extreme = int(np.sum(abs_weights > 10 * median_abs_w))
+    else:
+        # If median is 0, count all nonzero weights as extreme
+        n_extreme = int(np.sum(abs_weights > 0))
+
+    # Warning level
+    if ess_ratio < 0.2:
+        warning_level = 'severe'
+    elif ess_ratio < 0.5:
+        warning_level = 'caution'
+    else:
+        warning_level = 'ok'
+
+    result = {
+        'ess': float(ess),
+        'ess_ratio': float(ess_ratio),
+        'weight_max': float(weight_max),
+        'weight_min': float(weight_min),
+        'weight_ratio': float(weight_ratio) if np.isfinite(weight_ratio) else weight_ratio,
+        'cv': float(cv) if np.isfinite(cv) else cv,
+        'n_extreme': n_extreme,
+        'n_negative': n_negative,
+        'warning_level': warning_level,
+        'group_diagnostics': None,
+    }
+
+    # Group-specific diagnostics
+    if treat is not None:
+        treat = np.asarray(treat).ravel()
+        if len(treat) == n:
+            group_diag = {}
+            for level in np.unique(treat):
+                mask = treat == level
+                g_weights = weights[mask]
+                g_n = len(g_weights)
+
+                # Use absolute values for ESS when negatives present
+                g_abs_w = np.abs(g_weights)
+                g_sum_w = np.sum(g_abs_w)
+                g_sum_w2 = np.sum(g_abs_w ** 2)
+
+                if g_sum_w2 > 0:
+                    g_ess = (g_sum_w ** 2) / g_sum_w2
+                else:
+                    g_ess = 0.0
+
+                group_diag[level] = {
+                    'n': g_n,
+                    'ess': float(g_ess),
+                    'ess_ratio': float(g_ess / g_n) if g_n > 0 else 0.0,
+                    'weight_mean': float(np.mean(g_weights)),
+                    'weight_max': float(np.max(g_abs_w)) if g_n > 0 else np.nan,
+                    'n_negative': int(np.sum(g_weights < 0)),
+                }
+            result['group_diagnostics'] = group_diag
+
+    return result

cbps/highdim/__init__.py

@@ -0,0 +1,84 @@
+"""
+High-Dimensional Covariate Balancing Propensity Score (hdCBPS)
+==============================================================
+
+This module implements the High-Dimensional Covariate Balancing Propensity
+Score (hdCBPS) methodology for robust causal inference in settings where
+the number of covariates may exceed the sample size (p >> n).
+
+Algorithm Overview
+------------------
+The hdCBPS algorithm proceeds in four steps as described in Ning et al. (2020):
+
+1. **Propensity Score Estimation** (Equation 5): Fit penalized logistic
+   regression (LASSO) to obtain initial propensity score coefficients.
+
+2. **Outcome Model Estimation** (Equation 6): Fit penalized regression
+   (LASSO) to estimate outcome model coefficients separately for treatment
+   and control groups. This implementation uses unweighted LASSO (w_2=1).
+
+3. **Covariate Balancing** (Equation 7): Calibrate the propensity score by
+   minimizing the GMM objective to balance covariates selected in Step 2.
+   This achieves the weak covariate balancing property (Equation 9).
+
+4. **Treatment Effect Estimation**: Compute ATE/ATT using the Horvitz-Thompson
+   estimator with calibrated propensity scores. Standard errors are computed
+   using the sandwich variance estimator (Equation 11).
+
+Key Features
+------------
+- **Double Robustness**: Consistent and asymptotically normal when either
+  the propensity score model or outcome model is correctly specified.
+- **Sample Boundedness**: Estimated ATE lies within the range of observed
+  outcomes, ensuring stable estimates.
+- **Semiparametric Efficiency**: Achieves the efficiency bound when both
+  models are correctly specified.
+- **High-Dimensional Support**: Handles p >> n through L1 regularization.
+
+Requirements
+------------
+- **glmnetforpython**: Required for LASSO regularization with Fortran backend.
+- numpy, scipy: Numerical computations.
+- pandas: Data handling.
+
+References
+----------
+Ning, Y., Peng, S., and Imai, K. (2020). Robust estimation of causal effects
+via a high-dimensional covariate balancing propensity score. Biometrika,
+107(3), 533-554. https://doi.org/10.1093/biomet/asaa020
+
+See Also
+--------
+cbps.CBPS : Standard CBPS for low-dimensional settings.
+cbps.CBPSContinuous : CBPS for continuous treatments.
+"""
+
+__all__ = []
+
+# Import hdCBPS function when glmnet is available
+try:
+    from .hdcbps import hdCBPS, HDCBPSResults
+    from .lasso_utils import cv_glmnet, select_variables
+    __all__.extend(['hdCBPS', 'HDCBPSResults', 'cv_glmnet', 'select_variables'])
+except ImportError:
+    import warnings
+    warnings.warn(
+        "hdCBPS requires glmnetforpython. Install with: "
+        "pip install glmnetforpython",
+        ImportWarning
+    )
+
+# Weight functions (available regardless of glmnet)
+from .weight_funcs import (
+    ate_wt_func,
+    ate_wt_nl_func,
+    att_wt_func,
+    att_wt_nl_func
+)
+
+__all__.extend([
+    'ate_wt_func',
+    'ate_wt_nl_func',
+    'att_wt_func',
+    'att_wt_nl_func'
+])

cbps/highdim/gmm_loss.py

@@ -0,0 +1,340 @@
+"""
+GMM Loss Functions for High-Dimensional CBPS
+=============================================
+
+This module implements the Generalized Method of Moments (GMM) loss functions
+used in Step 3 of the hdCBPS algorithm for covariate balance calibration.
+
+The GMM objective minimizes the squared norm of the covariate balancing
+moment conditions (Equation 7 in Ning et al., 2020):
+
+.. math::
+
+    \\tilde{\\gamma} = \\arg\\min_{\\gamma} \\|g_n(\\gamma)\\|_2^2
+
+where the moment function is:
+
+.. math::
+
+    g_n(\\gamma) = \\sum_{i=1}^{n}
+    \\left( \\frac{T_i}{\\pi(\\gamma^T X_{i\\tilde{S}} +
+    \\hat{\\beta}_{\\tilde{S}^c}^T X_{i\\tilde{S}^c})} - 1 \\right) X_{i\\tilde{S}}
+
+Note: The paper defines g_n with a 1/n factor, but this implementation uses
+the sum (without 1/n) since minimizing ||g||^2 and ||g/n||^2 yield identical
+solutions. This calibration step removes bias from the penalized estimators
+and achieves the weak covariate balancing property (Equation 9).
+
+References
+----------
+Ning, Y., Peng, S., and Imai, K. (2020). Robust estimation of causal effects
+via a high-dimensional covariate balancing propensity score.
+Biometrika, 107(3), 533-554. https://doi.org/10.1093/biomet/asaa020
+"""
+
+import numpy as np
+from typing import Tuple
+
+from .weight_funcs import ate_wt_func, ate_wt_nl_func, att_wt_func, att_wt_nl_func
+
+
+def gmm_func(
+    beta_curr: np.ndarray,
+    S: np.ndarray,
+    tt: int,
+    X_gmm: np.ndarray,
+    method: str,
+    cov1_coef: np.ndarray,
+    cov0_coef: np.ndarray,
+    treat: np.ndarray,
+    beta_ini: np.ndarray
+) -> float:
+    """
+    Compute the GMM loss function for ATE estimation.
+
+    This function evaluates the covariate balancing objective for average
+    treatment effect (ATE) estimation in Step 3 of hdCBPS. It computes
+    the squared norm of the moment conditions that enforce balance between
+    treatment groups.
+
+    Parameters
+    ----------
+    beta_curr : np.ndarray
+        Current coefficient estimates being optimized. Shape depends on method:
+
+        - Linear method: shape ``(len(S),)`` without intercept
+        - Nonlinear methods: shape ``(len(S)+1,)`` with intercept
+
+    S : np.ndarray
+        Indices of LASSO-selected variables from the outcome model (0-based).
+        Corresponds to :math:`\\tilde{S}` in the paper.
+    tt : int
+        Treatment group indicator:
+
+        - 0: Optimize for control group (estimating :math:`\\mu_0`)
+        - 1: Optimize for treated group (estimating :math:`\\mu_1`)
+
+    X_gmm : np.ndarray, shape (n, p)
+        Covariate matrix without intercept column.
+    method : str
+        Outcome model specification:
+
+        - ``'linear'``: Gaussian outcome model
+        - ``'binomial'``: Logistic outcome model
+        - ``'poisson'``: Poisson outcome model
+
+    cov1_coef : np.ndarray, shape (p+1,)
+        Outcome model coefficients for the treated group.
+    cov0_coef : np.ndarray, shape (p+1,)
+        Outcome model coefficients for the control group.
+    treat : np.ndarray, shape (n,)
+        Binary treatment indicator (0/1).
+    beta_ini : np.ndarray, shape (p+1,)
+        Initial propensity score coefficients from LASSO (Step 1).
+
+    Returns
+    -------
+    loss : float
+        GMM loss value: :math:`\\|g_n(\\gamma)\\|_2^2`.
+
+    Notes
+    -----
+    For the linear outcome model, the moment condition is:
+
+    .. math::
+
+        g_n(\\gamma) = \\sum_{i=1}^{n}
+        \\left( \\frac{T_i}{\\pi_i} - 1 \\right) X_{i\\tilde{S}}
+
+    For generalized linear models (binomial/poisson), the weighted covariates
+    :math:`f(X) = b''(\\tilde{\\alpha}^T X) X_{\\tilde{S}}` are balanced instead,
+    as described in Section 4 of the paper.
+    """
+    # Convert covariate matrix to numpy array
+    x1 = np.asarray(X_gmm)
+    n1 = x1.shape[0]
+    
+    # IMPORTANT: Match R's behavior exactly
+    # R code does: X1 = cbind(rep(1,n1), x1) even though x1 already has intercept
+    # This adds an extra intercept column, making X1 have shape (n, p+2)
+    # The S indices from coef() are 1-based in R and refer to the (p+1) coefficient vector
+    # When used with X1[,S], the extra intercept shifts the column alignment
+    # We must replicate this behavior for compatibility
+    X1 = np.column_stack([np.ones(n1), x1])
+    
+    # Branch on method type
+    if method == "linear":
+        # Linear method: simple inverse probability weights
+        W = ate_wt_func(beta_curr, S, tt, x1, beta_ini, treat)
+        
+        # Compute weighted covariate means
+        if len(S) > 0:
+            # Extract selected columns and compute weighted means
+            w_curr_del = X1[:, S].T @ W
+            w_curr_del = np.asarray(w_curr_del).ravel()
+        else:
+            # No selected variables
+            w_curr_del = np.array([])
+    
+    elif method == "poisson":
+        # Poisson method: exponential link weights
+        W = ate_wt_nl_func(beta_curr, S, tt, x1, beta_ini, treat)
+        
+        # Select outcome model coefficients by treatment group
+        if tt == 1:
+            pweight = np.exp(X1 @ cov1_coef)
+        else:
+            pweight = np.exp(X1 @ cov0_coef)
+        
+        # Compute weighted covariates
+        if len(S) > 0:
+            # Stack outcome weights with weighted selected covariates
+            weighted_X = np.column_stack([
+                pweight,
+                pweight[:, None] * X1[:, S]
+            ])
+            w_curr_del = weighted_X.T @ W
+            w_curr_del = np.asarray(w_curr_del).ravel()
+        else:
+            # Only outcome weight when no covariates selected
+            w_curr_del = pweight @ W
+            w_curr_del = np.asarray([w_curr_del]).ravel()
+    
+    elif method == "binomial":
+        # Binomial method: logistic link weights
+        W = ate_wt_nl_func(beta_curr, S, tt, x1, beta_ini, treat)
+        
+        # Compute logistic probabilities and derivatives
+        if tt == 1:
+            # Treated group outcome model
+            exp_term = np.exp(X1 @ cov1_coef)
+            pweight1 = exp_term / (1.0 + exp_term)
+            pweight2 = exp_term / (1.0 + exp_term)**2
+        else:
+            # Control group outcome model
+            exp_term = np.exp(X1 @ cov0_coef)
+            pweight1 = exp_term / (1.0 + exp_term)
+            pweight2 = exp_term / (1.0 + exp_term)**2
+        
+        # Compute weighted covariates
+        if len(S) > 0:
+            # Stack probability with derivative-weighted selected covariates
+            weighted_X = np.column_stack([
+                pweight1,
+                pweight2[:, None] * X1[:, S]
+            ])
+            w_curr_del = weighted_X.T @ W
+            w_curr_del = np.asarray(w_curr_del).ravel()
+        else:
+            # Only probability weight when no covariates selected
+            w_curr_del = pweight1 @ W
+            w_curr_del = np.asarray([w_curr_del]).ravel()
+    
+    else:
+        raise ValueError(
+            f"method '{method}' not supported. "
+            f"Choose from: 'linear', 'binomial', 'poisson'"
+        )
+    
+    # Compute GMM loss as squared norm of moment conditions
+    gbar = w_curr_del
+    loss = gbar @ gbar
+    
+    return float(loss)
+
+
+def att_gmm_func(
+    beta_curr: np.ndarray,
+    S: np.ndarray,
+    X_gmm: np.ndarray,
+    method: str,
+    cov0_coef: np.ndarray,
+    treat: np.ndarray,
+    beta_ini: np.ndarray
+) -> float:
+    """
+    Compute the GMM loss function for ATT estimation.
+
+    This function evaluates the covariate balancing objective for the average
+    treatment effect on the treated (ATT) in Step 3 of hdCBPS. For ATT, only
+    the control group propensity score is calibrated to match the treated
+    group covariate distribution.
+
+    Parameters
+    ----------
+    beta_curr : np.ndarray
+        Current coefficient estimates being optimized. Shape depends on method:
+
+        - Linear method: shape ``(len(S),)`` without intercept
+        - Nonlinear methods: shape ``(len(S)+1,)`` with intercept
+
+    S : np.ndarray
+        Indices of LASSO-selected variables from the control outcome model.
+    X_gmm : np.ndarray, shape (n, p)
+        Covariate matrix without intercept column.
+    method : str
+        Outcome model specification: ``'linear'``, ``'binomial'``, or ``'poisson'``.
+    cov0_coef : np.ndarray, shape (p+1,)
+        Outcome model coefficients for the control group.
+    treat : np.ndarray, shape (n,)
+        Binary treatment indicator (0/1).
+    beta_ini : np.ndarray, shape (p+1,)
+        Initial propensity score coefficients from LASSO.
+
+    Returns
+    -------
+    loss : float
+        GMM loss value: :math:`\\|g_n(\\gamma)\\|_2^2`.
+
+    Notes
+    -----
+    Unlike ATE estimation which requires separate optimization for treated
+    and control groups, ATT estimation only requires calibrating the control
+    group weights to match the treated group distribution. The ATT moment
+    condition ensures that the reweighted control group has the same covariate
+    means as the treated group.
+
+    See the Supplementary Material of Ning et al. (2020) for theoretical
+    details on ATT estimation in high-dimensional settings.
+    """
+    # Convert covariate matrix to numpy array
+    x1 = np.asarray(X_gmm)
+    n1 = x1.shape[0]
+    
+    # IMPORTANT: Match R's behavior exactly
+    # R code does: X1 = cbind(rep(1,n1), x1) even though x1 already has intercept
+    # This adds an extra intercept column, making X1 have shape (n, p+2)
+    # The S indices from coef() are 1-based in R and refer to the (p+1) coefficient vector
+    # When used with X1[,S], the extra intercept shifts the column alignment
+    # We must replicate this behavior for compatibility
+    X1 = np.column_stack([np.ones(n1), x1])
+    
+    # Branch on method type
+    if method == "linear":
+        # Linear method: simple inverse probability weights for ATT
+        W = att_wt_func(beta_curr, S, x1, beta_ini, treat)
+        
+        # Compute weighted covariate means
+        if len(S) > 0:
+            w_curr_del = X1[:, S].T @ W
+            w_curr_del = np.asarray(w_curr_del).ravel()
+        else:
+            w_curr_del = np.array([])
+    
+    elif method == "poisson":
+        # Poisson method: exponential link weights for ATT
+        W = att_wt_nl_func(beta_curr, S, x1, beta_ini, treat)
+        
+        # Compute exponential weights from control outcome model
+        # Note: X1 = cbind(1, x1) has shape (n, p+2), cov0_coef has length p+2
+        pweight = np.exp(X1 @ cov0_coef)
+        
+        # Compute weighted covariates
+        if len(S) > 0:
+            # Stack outcome weights with weighted selected covariates
+            weighted_X = np.column_stack([
+                pweight,
+                pweight[:, None] * X1[:, S]
+            ])
+            w_curr_del = weighted_X.T @ W
+            w_curr_del = np.asarray(w_curr_del).ravel()
+        else:
+            # Only outcome weight when no covariates selected
+            w_curr_del = pweight @ W
+            w_curr_del = np.asarray([w_curr_del]).ravel()
+    
+    elif method == "binomial":
+        # Binomial method: logistic link weights for ATT
+        W = att_wt_nl_func(beta_curr, S, x1, beta_ini, treat)
+        
+        # Compute logistic probabilities and derivatives (control group only)
+        # Note: X1 = cbind(1, x1) has shape (n, p+2), cov0_coef has length p+2
+        exp_term = np.exp(X1 @ cov0_coef)
+        pweight1 = exp_term / (1.0 + exp_term)
+        pweight2 = exp_term / (1.0 + exp_term)**2
+        
+        # Compute weighted covariates
+        if len(S) > 0:
+            # Stack probability with derivative-weighted selected covariates
+            weighted_X = np.column_stack([
+                pweight1,
+                pweight2[:, None] * X1[:, S]
+            ])
+            w_curr_del = weighted_X.T @ W
+            w_curr_del = np.asarray(w_curr_del).ravel()
+        else:
+            # Only probability weight when no covariates selected
+            w_curr_del = pweight1 @ W
+            w_curr_del = np.asarray([w_curr_del]).ravel()
+    
+    else:
+        raise ValueError(
+            f"method '{method}' not supported. "
+            f"Choose from: 'linear', 'binomial', 'poisson'"
+        )
+    
+    # Compute GMM loss as squared norm of moment conditions
+    gbar = w_curr_del
+    loss = gbar @ gbar
+    
+    return float(loss)