cbps 0.2.0__py3-none-any.whl

cbps/diagnostics/balance_cbmsm_addon.py

@@ -0,0 +1,162 @@
+"""
+Covariate Balance for Marginal Structural Models
+=================================================
+
+Balance diagnostics for marginal structural models estimated via CBMSM.
+These diagnostics assess whether CBMSM weights achieve covariate balance
+across treatment history groups in longitudinal settings with time-varying
+confounding.
+
+The balance_cbmsm function computes weighted covariate means within each
+treatment trajectory pattern (e.g., "0+1+1" for control-treated-treated),
+comparing CBMSM weights against standard GLM-based inverse probability weights.
+
+References
+----------
+Imai, K. and Ratkovic, M. (2015). Robust estimation of inverse probability
+weights for marginal structural models. Journal of the American Statistical
+Association, 110(511), 1013-1023.
+"""
+import numpy as np
+from typing import Dict, Any
+
+
+def balance_cbmsm(cbmsm_obj: Dict[str, Any]) -> Dict[str, np.ndarray]:
+    """
+    Compute covariate balance across treatment history groups.
+
+    Calculates weighted covariate means within each treatment trajectory
+    pattern, enabling assessment of balance in marginal structural models.
+    Treatment histories are represented as concatenated strings (e.g.,
+    "0+1+1" for a unit that was untreated in period 1, then treated in
+    periods 2 and 3).
+
+    Parameters
+    ----------
+    cbmsm_obj : dict
+        Fitted CBMSM result containing:
+
+        - **y** : ndarray of shape (n_obs,) - Binary treatment at each observation
+        - **x** : ndarray of shape (n_obs, k) - Covariates with intercept
+        - **weights** : ndarray - CBMSM weights (unit-level)
+        - **id** : ndarray - Unit identifiers
+        - **time** : ndarray - Time period indicators
+        - **glm_weights** : ndarray - Baseline GLM inverse probability weights
+
+    Returns
+    -------
+    dict
+        Balance statistics with keys:
+
+        - **Balanced** : ndarray of shape (n_covars, 2*n_hist)
+          CBMSM-weighted means (first n_hist columns) and standardized
+          means (remaining columns) for each treatment history.
+        - **Unweighted** : ndarray of shape (n_covars, 2*n_hist)
+          GLM-weighted baseline statistics for comparison.
+        - **StatBal** : float
+          Summary statistic measuring overall imbalance.
+
+    Notes
+    -----
+    Balance is assessed using first-period covariate values, following
+    standard practice for evaluating baseline balance in MSM settings.
+
+    See Also
+    --------
+    plot_cbmsm : Visualization of CBMSM balance diagnostics.
+    """
+    # Extract panel data components
+    y = cbmsm_obj['y']
+    x = cbmsm_obj['x']
+    w = cbmsm_obj.get('weights', cbmsm_obj.get('w'))
+    glm_w = cbmsm_obj.get('glm_weights', cbmsm_obj.get('glm.w', cbmsm_obj.get('glm_w')))
+    ids = cbmsm_obj['id']
+    times = cbmsm_obj['time']
+
+    unique_ids = np.sort(np.unique(ids))
+    unique_times = np.sort(np.unique(times))
+    n_units = len(unique_ids)
+    n_periods = len(unique_times)
+
+    # Reconstruct treatment history matrix (n_units x n_periods)
+    treat_hist = np.full((n_units, n_periods), np.nan)
+    for i, unit_id in enumerate(unique_ids):
+        for j, period in enumerate(unique_times):
+            mask = (ids == unit_id) & (times == period)
+            if mask.any():
+                treat_hist[i, j] = y[mask][0]
+
+    # Create treatment history factor (e.g., "0+1+1" for trajectory)
+    treat_hist_fac = []
+    for i in range(n_units):
+        hist_str = '+'.join(
+            [str(int(t)) if not np.isnan(t) else 'NA' for t in treat_hist[i, :]]
+        )
+        treat_hist_fac.append(hist_str)
+    treat_hist_fac = np.array(treat_hist_fac)
+
+    unique_hist = np.unique(treat_hist_fac)
+    n_hist = len(unique_hist)
+
+    # Initialize balance matrices
+    n_covars = x.shape[1] - 1  # Exclude intercept column
+    bal = np.zeros((n_covars, 2 * n_hist))
+    baseline = np.zeros((n_covars, 2 * n_hist))
+
+    # Use first-period covariates for balance assessment
+    first_period = unique_times[0]
+    first_period_mask = (times == first_period)
+    x_first = x[first_period_mask, :]
+    ids_first = ids[first_period_mask]
+
+    # Map unit-level weights to first-period observations
+    if len(w) == n_units:
+        id_to_w = dict(zip(unique_ids, w))
+        id_to_glm_w = dict(zip(unique_ids, glm_w))
+        w_first = np.array([id_to_w[uid] for uid in ids_first])
+        glm_w_first = np.array([id_to_glm_w[uid] for uid in ids_first])
+    else:
+        w_first = w[first_period_mask]
+        glm_w_first = glm_w[first_period_mask]
+
+    # Compute balance statistics for each treatment history group
+    for i, hist in enumerate(unique_hist):
+        hist_mask_units = (treat_hist_fac == hist)
+        hist_mask_obs = np.array(
+            [uid in unique_ids[hist_mask_units] for uid in ids_first]
+        )
+
+        # Iterate over covariates (skip intercept at index 0)
+        for j in range(1, x.shape[1]):
+            idx = j - 1
+
+            if hist_mask_obs.sum() > 0:
+                # CBMSM-weighted mean
+                bal[idx, i] = (
+                    np.sum(hist_mask_obs * x_first[:, j] * w_first) /
+                    np.sum(w_first * hist_mask_obs)
+                )
+                # Standardized mean (divided by weighted std)
+                bal[idx, i + n_hist] = bal[idx, i] / np.std(w_first * x_first[:, j])
+
+                # GLM baseline weighted mean
+                baseline[idx, i] = (
+                    np.sum(hist_mask_obs * x_first[:, j] * glm_w_first) /
+                    np.sum(glm_w_first * hist_mask_obs)
+                )
+                baseline[idx, i + n_hist] = (
+                    baseline[idx, i] / np.std(glm_w_first * x_first[:, j])
+                )
+
+    # Handle numerical edge cases
+    bal[np.isnan(bal)] = 0
+    baseline[np.isnan(baseline)] = 0
+
+    # Summary balance statistic
+    statbal = np.sum((bal - bal[:, 0:1]) * (bal != 0) ** 2)
+
+    return {
+        'Balanced': bal,
+        'Unweighted': baseline,
+        'StatBal': statbal
+    }

cbps/diagnostics/continuous_diagnostics.py

@@ -0,0 +1,259 @@
+"""
+Continuous Treatment CBPS Diagnostic Tools
+==========================================
+
+Diagnostic utilities for assessing the quality of CBGPS (Covariate Balancing
+Generalized Propensity Score) estimation results.
+
+Author: CBPS Development Team
+Date: 2026-01-28
+"""
+
+import numpy as np
+import warnings
+from typing import Dict, Tuple, Optional
+import statsmodels.api as sm
+
+
+def compute_f_statistic(treat: np.ndarray, X: np.ndarray, weights: np.ndarray) -> Tuple[float, float]:
+    """
+    Compute F-statistic from weighted regression of treatment on covariates.
+    
+    This is the primary balance metric used in Fong et al. (2018).
+    Lower F-statistic indicates better covariate balance.
+    
+    Parameters
+    ----------
+    treat : np.ndarray
+        Treatment variable (n,)
+    X : np.ndarray
+        Covariate matrix (n, K), should NOT include intercept
+    weights : np.ndarray
+        CBGPS weights (n,)
+        
+    Returns
+    -------
+    f_stat : float
+        F-statistic from weighted regression T ~ X
+    p_value : float
+        P-value for the F-test
+    """
+    n, K = X.shape
+    X_with_const = sm.add_constant(X)
+    
+    # Normalize weights
+    weights_norm = weights * n / weights.sum()
+    
+    # Weighted least squares
+    model = sm.WLS(treat, X_with_const, weights=weights_norm)
+    result = model.fit()
+    
+    return result.fvalue, result.f_pvalue
+
+
+def compute_weighted_correlations(treat: np.ndarray, X: np.ndarray, 
+                                  weights: np.ndarray) -> np.ndarray:
+    """
+    Compute weighted Pearson correlations between treatment and each covariate.
+    
+    Parameters
+    ----------
+    treat : np.ndarray
+        Treatment variable (n,)
+    X : np.ndarray
+        Covariate matrix (n, K)
+    weights : np.ndarray
+        CBGPS weights (n,)
+        
+    Returns
+    -------
+    np.ndarray
+        Weighted correlations for each covariate (K,)
+    """
+    n, K = X.shape
+    w_sum = np.sum(weights)
+    w_mean_T = np.sum(weights * treat) / w_sum
+    
+    correlations = []
+    for j in range(K):
+        X_j = X[:, j]
+        w_mean_X = np.sum(weights * X_j) / w_sum
+        
+        w_cov = np.sum(weights * (treat - w_mean_T) * (X_j - w_mean_X)) / w_sum
+        w_var_T = np.sum(weights * (treat - w_mean_T)**2) / w_sum
+        w_var_X = np.sum(weights * (X_j - w_mean_X)**2) / w_sum
+        
+        if w_var_T > 1e-10 and w_var_X > 1e-10:
+            corr = w_cov / np.sqrt(w_var_T * w_var_X)
+        else:
+            corr = np.nan
+        
+        correlations.append(corr)
+    
+    return np.array(correlations)
+
+
+def diagnose_cbgps_quality(cbgps_result: Dict, treat: np.ndarray, 
+                           X: np.ndarray) -> Dict[str, any]:
+    """
+    Diagnose the quality of CBGPS estimation results.
+    
+    This function assesses covariate balance quality and provides
+    recommendations for users.
+    
+    Parameters
+    ----------
+    cbgps_result : dict
+        Result dictionary from cbps_continuous_fit()
+    treat : np.ndarray
+        Treatment variable (n,)
+    X : np.ndarray
+        Covariate matrix (n, K), without intercept
+        
+    Returns
+    -------
+    dict
+        Diagnostic information including:
+        - quality_level: 'excellent', 'acceptable', 'moderate', or 'poor'
+        - f_statistic: Overall balance F-statistic
+        - max_abs_correlation: Maximum absolute weighted correlation
+        - mean_abs_correlation: Mean absolute weighted correlation
+        - recommendation: User-facing recommendation string
+        - j_statistic: GMM objective value (if available)
+        
+    Examples
+    --------
+    >>> from cbps.core import cbps_continuous_fit
+    >>> result = cbps_continuous_fit(treat, X, method='exact')
+    >>> diag = diagnose_cbgps_quality(result, treat, X)
+    >>> print(f"Balance quality: {diag['quality_level']}")
+    >>> print(f"Recommendation: {diag['recommendation']}")
+    """
+    weights = cbgps_result.get('weights')
+    if weights is None:
+        raise ValueError("cbgps_result must contain 'weights' key")
+    
+    # Compute balance metrics
+    f_stat, f_pval = compute_f_statistic(treat, X, weights)
+    correlations = compute_weighted_correlations(treat, X, weights)
+    max_abs_corr = np.max(np.abs(correlations))
+    mean_abs_corr = np.mean(np.abs(correlations))
+    
+    # Determine quality level
+    if f_stat < 0.1 and max_abs_corr < 0.05:
+        quality_level = 'excellent'
+        recommendation = (
+            "Excellent covariate balance achieved. "
+            "Proceed with confidence using these weights."
+        )
+    elif f_stat < 0.5 and max_abs_corr < 0.15:
+        quality_level = 'acceptable'
+        recommendation = (
+            "Acceptable covariate balance. "
+            "Check balance_table() for any concerning imbalances."
+        )
+    elif f_stat < 2.0:
+        quality_level = 'moderate'
+        recommendation = (
+            "Moderate balance quality. Consider: "
+            "(1) using npCBPS for more robust results, "
+            "(2) reducing the number of covariates, or "
+            "(3) increasing sample size if possible."
+        )
+    else:
+        quality_level = 'poor'
+        recommendation = (
+            "Poor balance quality detected. Strongly recommend: "
+            "(1) switching to npCBPS (method='nonparametric'), "
+            "(2) careful variable selection to reduce noise covariates, or "
+            "(3) verifying data quality and model specification."
+        )
+    
+    # Collect diagnostic info
+    diagnostic = {
+        'quality_level': quality_level,
+        'f_statistic': float(f_stat),
+        'f_pvalue': float(f_pval),
+        'max_abs_correlation': float(max_abs_corr),
+        'mean_abs_correlation': float(mean_abs_corr),
+        'correlations': correlations,
+        'recommendation': recommendation,
+        'converged': cbgps_result.get('converged', None),
+        'j_statistic': cbgps_result.get('J', None),
+        'weight_range': (float(np.min(weights)), float(np.max(weights))),
+        'weight_sum': float(np.sum(weights)),
+    }
+    
+    return diagnostic
+
+
+def print_balance_diagnosis(diagnostic: Dict) -> None:
+    """
+    Print a user-friendly diagnostic report.
+    
+    Parameters
+    ----------
+    diagnostic : dict
+        Output from diagnose_cbgps_quality()
+    """
+    print("=" * 70)
+    print("CBGPS Balance Quality Diagnostic Report")
+    print("=" * 70)
+    print()
+    
+    # Quality badge
+    quality = diagnostic['quality_level']
+    badges = {
+        'excellent': '✓✓✓ EXCELLENT',
+        'acceptable': '✓✓ ACCEPTABLE',
+        'moderate': '⚠ MODERATE',
+        'poor': '✗ POOR'
+    }
+    print(f"Overall Quality: {badges.get(quality, quality.upper())}")
+    print()
+    
+    # Metrics
+    print("Balance Metrics:")
+    print(f"  F-statistic: {diagnostic['f_statistic']:.4f} (p={diagnostic['f_pvalue']:.4f})")
+    print(f"  Max |correlation|: {diagnostic['max_abs_correlation']:.4f}")
+    print(f"  Mean |correlation|: {diagnostic['mean_abs_correlation']:.4f}")
+    
+    if diagnostic['j_statistic'] is not None:
+        print(f"  J-statistic: {diagnostic['j_statistic']:.6f}")
+    
+    if diagnostic['converged'] is not None:
+        conv_str = "Yes" if diagnostic['converged'] else "No"
+        print(f"  Converged: {conv_str}")
+    
+    print()
+    
+    # Weight diagnostics
+    w_min, w_max = diagnostic['weight_range']
+    print(f"Weight Diagnostics:")
+    print(f"  Sum of weights: {diagnostic['weight_sum']:.4f} (should be ~1.0)")
+    print(f"  Weight range: [{w_min:.4f}, {w_max:.4f}]")
+    
+    if w_max / w_min > 100:
+        print("  ⚠ Warning: Extreme weight ratio detected")
+    
+    print()
+    
+    # Recommendation
+    print("Recommendation:")
+    print(f"  {diagnostic['recommendation']}")
+    print()
+    print("=" * 70)
+
+
+# Example usage
+if __name__ == "__main__":
+    # This is a demonstration - normally called after cbps_continuous_fit()
+    print("Continuous CBPS Diagnostic Tools")
+    print()
+    print("Usage:")
+    print("  from cbps.diagnostics.continuous_diagnostics import diagnose_cbgps_quality")
+    print("  from cbps.core import cbps_continuous_fit")
+    print()
+    print("  result = cbps_continuous_fit(treat, X, method='exact')")
+    print("  diag = diagnose_cbgps_quality(result, treat, X)")
+    print("  print_balance_diagnosis(diag)")

cbps/diagnostics/normality.py

@@ -0,0 +1,173 @@
+"""Normality diagnostics for continuous treatment CBPS.
+
+Fong, Hazlett, and Imai (2018) assume the treatment variable is
+conditionally normal: T | X ~ N(X'beta, sigma^2). This module provides
+diagnostic tools to check this assumption.
+
+References
+----------
+Fong, C., Hazlett, C., and Imai, K. (2018). Covariate balancing propensity
+score for a continuous treatment. The Annals of Applied Statistics, 12(1):
+156-177.
+"""
+
+import numpy as np
+from scipy import stats
+from typing import Dict, Any
+import warnings
+
+
+def test_treatment_normality(
+    treat: np.ndarray,
+    X: np.ndarray,
+    alpha: float = 0.05,
+) -> Dict[str, Any]:
+    """Test normality of treatment residuals (T - X'beta_hat).
+
+    Under the continuous CBPS model, T | X ~ N(X'beta, sigma^2). This
+    function fits OLS to obtain residuals e = T - X @ beta_hat, then
+    tests whether these residuals follow a normal distribution.
+
+    For n <= 5000, the Shapiro-Wilk test is used (highest power for
+    detecting departures from normality in moderate samples). For n > 5000,
+    the D'Agostino-Pearson omnibus test is used (Shapiro-Wilk becomes
+    computationally expensive and overly sensitive for very large n).
+
+    Parameters
+    ----------
+    treat : np.ndarray, shape (n,)
+        Continuous treatment variable.
+    X : np.ndarray, shape (n, k)
+        Covariate matrix (should include intercept if desired).
+    alpha : float, default=0.05
+        Significance level for the normality test.
+
+    Returns
+    -------
+    dict with:
+        - statistic : float
+            Test statistic value.
+        - p_value : float
+            p-value of the normality test.
+        - test_used : str
+            Name of the test applied ('shapiro-wilk' or 'dagostino-pearson').
+        - reject_normality : bool
+            True if normality is rejected at the given alpha level.
+        - skewness : float
+            Sample skewness of residuals.
+        - kurtosis : float
+            Sample excess kurtosis of residuals.
+        - warning_message : str or None
+            If normality is rejected, suggests using npCBPS.
+
+    Notes
+    -----
+    If normality is rejected, the parametric CBPS estimator may be
+    misspecified. Consider using the nonparametric variant (npCBPS)
+    which does not require distributional assumptions on T | X.
+    """
+    treat = np.asarray(treat, dtype=float).ravel()
+    X = np.asarray(X, dtype=float)
+    n = len(treat)
+
+    if X.ndim == 1:
+        X = X.reshape(-1, 1)
+
+    # Input validation: NaN/Inf checks
+    if np.any(~np.isfinite(treat)):
+        return {
+            'statistic': np.nan, 'p_value': np.nan, 'test_used': 'none',
+            'reject_normality': False, 'skewness': np.nan, 'kurtosis': np.nan,
+            'n_observations': n,
+            'warning_message': "Treatment variable contains NaN or Inf values. "
+                              "Cannot perform normality test. Clean data first."
+        }
+    if np.any(~np.isfinite(X)):
+        return {
+            'statistic': np.nan, 'p_value': np.nan, 'test_used': 'none',
+            'reject_normality': False, 'skewness': np.nan, 'kurtosis': np.nan,
+            'n_observations': n,
+            'warning_message': "Covariate matrix contains NaN or Inf values. "
+                              "Cannot perform normality test. Clean data first."
+        }
+
+    # Minimum sample size check
+    if n < 3:
+        return {
+            'statistic': np.nan, 'p_value': np.nan, 'test_used': 'none',
+            'reject_normality': False, 'skewness': np.nan, 'kurtosis': np.nan,
+            'n_observations': n,
+            'warning_message': f"Sample size n={n} too small for normality testing "
+                              f"(minimum 3 required for Shapiro-Wilk)."
+        }
+
+    if X.shape[0] != n:
+        raise ValueError(
+            f"Dimension mismatch: treat has {n} observations but X has "
+            f"{X.shape[0]} rows."
+        )
+
+    # Fit OLS: beta_hat = (X'X)^{-1} X'T
+    # Use lstsq for numerical stability
+    beta_hat, _, _, _ = np.linalg.lstsq(X, treat, rcond=None)
+    residuals = treat - X @ beta_hat
+
+    # Near-zero residual variance guard
+    resid_std = np.std(residuals, ddof=1)
+    if resid_std < 1e-10:
+        return {
+            'statistic': np.nan, 'p_value': np.nan, 'test_used': 'none',
+            'reject_normality': False, 'skewness': 0.0, 'kurtosis': 0.0,
+            'n_observations': n,
+            'warning_message': "Treatment is a near-perfect linear function of covariates "
+                              f"(residual std = {resid_std:.2e}). Normality test not applicable."
+        }
+
+    # Compute descriptive statistics
+    skewness = float(stats.skew(residuals))
+    kurtosis = float(stats.kurtosis(residuals))  # excess kurtosis
+
+    # Choose test based on sample size
+    if n <= 5000:
+        test_used = "shapiro-wilk"
+        stat, p_value = stats.shapiro(residuals)
+    else:
+        test_used = "dagostino-pearson"
+        stat, p_value = stats.normaltest(residuals)
+
+    stat = float(stat)
+    p_value = float(p_value)
+
+    # NaN result guard
+    if not np.isfinite(stat) or not np.isfinite(p_value):
+        return {
+            'statistic': stat, 'p_value': p_value, 'test_used': test_used,
+            'reject_normality': False, 'skewness': np.nan, 'kurtosis': np.nan,
+            'n_observations': n,
+            'warning_message': f"Normality test ({test_used}) returned non-finite result. "
+                              f"This may indicate degenerate data or numerical issues."
+        }
+
+    reject = p_value < alpha
+
+    # Construct warning message
+    warning_message = None
+    if reject:
+        warning_message = (
+            f"Normality of treatment residuals rejected ({test_used} test, "
+            f"p={p_value:.4g}, alpha={alpha}). The conditional normality "
+            f"assumption T|X ~ N(X'beta, sigma^2) may not hold. Consider "
+            f"using npCBPS (nonparametric CBPS) which does not require "
+            f"distributional assumptions."
+        )
+        warnings.warn(warning_message, UserWarning, stacklevel=2)
+
+    return {
+        "statistic": stat,
+        "p_value": p_value,
+        "test_used": test_used,
+        "reject_normality": reject,
+        "skewness": skewness,
+        "kurtosis": kurtosis,
+        "warning_message": warning_message,
+    }