""" Comprehensive statistical assumption checking utilities. This module provides functions to check common statistical assumptions: - Normality - Homogeneity of variance - Independence - Linearity - Outliers """ import numpy as np import pandas as pd from scipy import stats import matplotlib.pyplot as plt import seaborn as sns from typing import Dict, List, Tuple, Optional, Union def check_normality( data: Union[np.ndarray, pd.Series, List], name: str = "data", alpha: float = 0.05, plot: bool = True ) -> Dict: """ Check normality assumption using Shapiro-Wilk test and visualizations. Parameters ---------- data : array-like Data to check for normality name : str Name of the variable (for labeling) alpha : float Significance level for Shapiro-Wilk test plot : bool Whether to create Q-Q plot and histogram Returns ------- dict Results including test statistic, p-value, and interpretation """ data = np.asarray(data) data_clean = data[~np.isnan(data)] # Shapiro-Wilk test statistic, p_value = stats.shapiro(data_clean) # Interpretation is_normal = p_value > alpha interpretation = ( f"Data {'appear' if is_normal else 'do not appear'} normally distributed " f"(W = {statistic:.3f}, p = {p_value:.3f})" ) # Visual checks if plot: fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) # Q-Q plot stats.probplot(data_clean, dist="norm", plot=ax1) ax1.set_title(f"Q-Q Plot: {name}") ax1.grid(alpha=0.3) # Histogram with normal curve ax2.hist(data_clean, bins='auto', density=True, alpha=0.7, color='steelblue', edgecolor='black') mu, sigma = data_clean.mean(), data_clean.std() x = np.linspace(data_clean.min(), data_clean.max(), 100) ax2.plot(x, stats.norm.pdf(x, mu, sigma), 'r-', linewidth=2, label='Normal curve') ax2.set_xlabel('Value') ax2.set_ylabel('Density') ax2.set_title(f'Histogram: {name}') ax2.legend() ax2.grid(alpha=0.3) plt.tight_layout() plt.show() return { 'test': 'Shapiro-Wilk', 'statistic': statistic, 'p_value': p_value, 'is_normal': is_normal, 'interpretation': interpretation, 'n': len(data_clean), 'recommendation': ( "Proceed with parametric test" if is_normal else "Consider non-parametric alternative or transformation" ) } def check_normality_per_group( data: pd.DataFrame, value_col: str, group_col: str, alpha: float = 0.05, plot: bool = True ) -> pd.DataFrame: """ Check normality assumption for each group separately. Parameters ---------- data : pd.DataFrame Data containing values and group labels value_col : str Column name for values to check group_col : str Column name for group labels alpha : float Significance level plot : bool Whether to create Q-Q plots for each group Returns ------- pd.DataFrame Results for each group """ groups = data[group_col].unique() results = [] if plot: n_groups = len(groups) fig, axes = plt.subplots(1, n_groups, figsize=(5 * n_groups, 4)) if n_groups == 1: axes = [axes] for idx, group in enumerate(groups): group_data = data[data[group_col] == group][value_col].dropna() stat, p = stats.shapiro(group_data) results.append({ 'Group': group, 'N': len(group_data), 'W': stat, 'p-value': p, 'Normal': 'Yes' if p > alpha else 'No' }) if plot: stats.probplot(group_data, dist="norm", plot=axes[idx]) axes[idx].set_title(f"Q-Q Plot: {group}") axes[idx].grid(alpha=0.3) if plot: plt.tight_layout() plt.show() return pd.DataFrame(results) def check_homogeneity_of_variance( data: pd.DataFrame, value_col: str, group_col: str, alpha: float = 0.05, plot: bool = True ) -> Dict: """ Check homogeneity of variance using Levene's test. Parameters ---------- data : pd.DataFrame Data containing values and group labels value_col : str Column name for values group_col : str Column name for group labels alpha : float Significance level plot : bool Whether to create box plots Returns ------- dict Results including test statistic, p-value, and interpretation """ groups = [group[value_col].values for name, group in data.groupby(group_col)] # Levene's test (robust to non-normality) statistic, p_value = stats.levene(*groups) # Variance ratio (max/min) variances = [np.var(g, ddof=1) for g in groups] var_ratio = max(variances) / min(variances) is_homogeneous = p_value > alpha interpretation = ( f"Variances {'appear' if is_homogeneous else 'do not appear'} homogeneous " f"(F = {statistic:.3f}, p = {p_value:.3f}, variance ratio = {var_ratio:.2f})" ) if plot: fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) # Box plot data.boxplot(column=value_col, by=group_col, ax=ax1) ax1.set_title('Box Plots by Group') ax1.set_xlabel(group_col) ax1.set_ylabel(value_col) plt.sca(ax1) plt.xticks(rotation=45) # Variance plot group_names = data[group_col].unique() ax2.bar(range(len(variances)), variances, color='steelblue', edgecolor='black') ax2.set_xticks(range(len(variances))) ax2.set_xticklabels(group_names, rotation=45) ax2.set_ylabel('Variance') ax2.set_title('Variance by Group') ax2.grid(alpha=0.3, axis='y') plt.tight_layout() plt.show() return { 'test': 'Levene', 'statistic': statistic, 'p_value': p_value, 'is_homogeneous': is_homogeneous, 'variance_ratio': var_ratio, 'interpretation': interpretation, 'recommendation': ( "Proceed with standard test" if is_homogeneous else "Consider Welch's correction or transformation" ) } def check_linearity( x: Union[np.ndarray, pd.Series], y: Union[np.ndarray, pd.Series], x_name: str = "X", y_name: str = "Y" ) -> Dict: """ Check linearity assumption for regression. Parameters ---------- x : array-like Predictor variable y : array-like Outcome variable x_name : str Name of predictor y_name : str Name of outcome Returns ------- dict Visualization and recommendations """ x = np.asarray(x) y = np.asarray(y) # Fit linear regression slope, intercept, r_value, p_value, std_err = stats.linregress(x, y) y_pred = intercept + slope * x # Calculate residuals residuals = y - y_pred # Visualization fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) # Scatter plot with regression line ax1.scatter(x, y, alpha=0.6, s=50, edgecolors='black', linewidths=0.5) ax1.plot(x, y_pred, 'r-', linewidth=2, label=f'y = {intercept:.2f} + {slope:.2f}x') ax1.set_xlabel(x_name) ax1.set_ylabel(y_name) ax1.set_title('Scatter Plot with Regression Line') ax1.legend() ax1.grid(alpha=0.3) # Residuals vs fitted ax2.scatter(y_pred, residuals, alpha=0.6, s=50, edgecolors='black', linewidths=0.5) ax2.axhline(y=0, color='r', linestyle='--', linewidth=2) ax2.set_xlabel('Fitted values') ax2.set_ylabel('Residuals') ax2.set_title('Residuals vs Fitted Values') ax2.grid(alpha=0.3) plt.tight_layout() plt.show() return { 'r': r_value, 'r_squared': r_value ** 2, 'interpretation': ( "Examine residual plot. Points should be randomly scattered around zero. " "Patterns (curves, funnels) suggest non-linearity or heteroscedasticity." ), 'recommendation': ( "If non-linear pattern detected: Consider polynomial terms, " "transformations, or non-linear models" ) } def detect_outliers( data: Union[np.ndarray, pd.Series, List], name: str = "data", method: str = "iqr", threshold: float = 1.5, plot: bool = True ) -> Dict: """ Detect outliers using IQR method or z-score method. Parameters ---------- data : array-like Data to check for outliers name : str Name of variable method : str Method to use: 'iqr' or 'zscore' threshold : float Threshold for outlier detection For IQR: typically 1.5 (mild) or 3 (extreme) For z-score: typically 3 plot : bool Whether to create visualizations Returns ------- dict Outlier indices, values, and visualizations """ data = np.asarray(data) data_clean = data[~np.isnan(data)] if method == "iqr": q1 = np.percentile(data_clean, 25) q3 = np.percentile(data_clean, 75) iqr = q3 - q1 lower_bound = q1 - threshold * iqr upper_bound = q3 + threshold * iqr outlier_mask = (data_clean < lower_bound) | (data_clean > upper_bound) elif method == "zscore": z_scores = np.abs(stats.zscore(data_clean)) outlier_mask = z_scores > threshold lower_bound = data_clean.mean() - threshold * data_clean.std() upper_bound = data_clean.mean() + threshold * data_clean.std() else: raise ValueError("method must be 'iqr' or 'zscore'") outlier_indices = np.where(outlier_mask)[0] outlier_values = data_clean[outlier_mask] n_outliers = len(outlier_indices) pct_outliers = (n_outliers / len(data_clean)) * 100 if plot: fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) # Box plot bp = ax1.boxplot(data_clean, vert=True, patch_artist=True) bp['boxes'][0].set_facecolor('steelblue') ax1.set_ylabel('Value') ax1.set_title(f'Box Plot: {name}') ax1.grid(alpha=0.3, axis='y') # Scatter plot highlighting outliers x_coords = np.arange(len(data_clean)) ax2.scatter(x_coords[~outlier_mask], data_clean[~outlier_mask], alpha=0.6, s=50, color='steelblue', label='Normal', edgecolors='black', linewidths=0.5) if n_outliers > 0: ax2.scatter(x_coords[outlier_mask], data_clean[outlier_mask], alpha=0.8, s=100, color='red', label='Outliers', marker='D', edgecolors='black', linewidths=0.5) ax2.axhline(y=lower_bound, color='orange', linestyle='--', linewidth=1.5, label='Bounds') ax2.axhline(y=upper_bound, color='orange', linestyle='--', linewidth=1.5) ax2.set_xlabel('Index') ax2.set_ylabel('Value') ax2.set_title(f'Outlier Detection: {name}') ax2.legend() ax2.grid(alpha=0.3) plt.tight_layout() plt.show() return { 'method': method, 'threshold': threshold, 'n_outliers': n_outliers, 'pct_outliers': pct_outliers, 'outlier_indices': outlier_indices, 'outlier_values': outlier_values, 'lower_bound': lower_bound, 'upper_bound': upper_bound, 'interpretation': f"Found {n_outliers} outliers ({pct_outliers:.1f}% of data)", 'recommendation': ( "Investigate outliers for data entry errors. " "Consider: (1) removing if errors, (2) winsorizing, " "(3) keeping if legitimate, (4) using robust methods" ) } def check_regression_diagnostics( model, alpha: float = 0.05, plot: bool = True ) -> Dict: """ Run standard diagnostics on a fitted statsmodels OLS model. Produces the classic 4-panel diagnostic plot (residuals vs fitted, Q-Q, scale-location, residual histogram) plus formal tests: normality of residuals (Shapiro-Wilk), heteroscedasticity (Breusch-Pagan), autocorrelation (Durbin-Watson), and multicollinearity (VIF per predictor). Parameters ---------- model : statsmodels regression results A fitted model from ``sm.OLS(y, X).fit()`` where X includes a constant (``sm.add_constant``). alpha : float Significance level for the formal tests. plot : bool Whether to create the 4-panel diagnostic figure. Returns ------- dict Test results, VIF table, and per-assumption interpretations. """ from statsmodels.stats.diagnostic import het_breuschpagan from statsmodels.stats.stattools import durbin_watson from statsmodels.stats.outliers_influence import variance_inflation_factor residuals = np.asarray(model.resid) fitted = np.asarray(model.fittedvalues) exog = model.model.exog exog_names = list(model.model.exog_names) # Normality of residuals sw_stat, sw_p = stats.shapiro(residuals) # Heteroscedasticity (Breusch-Pagan) bp_lm, bp_p, _, _ = het_breuschpagan(residuals, exog) # Autocorrelation (Durbin-Watson: ~2 = none, <1.5 or >2.5 = concern) dw = durbin_watson(residuals) # Multicollinearity (VIF, skipping the constant) vif_rows = [] for i, name in enumerate(exog_names): if name.lower() in ("const", "intercept"): continue vif_rows.append({'Variable': name, 'VIF': variance_inflation_factor(exog, i)}) vif_table = pd.DataFrame(vif_rows) max_vif = vif_table['VIF'].max() if not vif_table.empty else np.nan if plot: fig, axes = plt.subplots(2, 2, figsize=(12, 10)) axes[0, 0].scatter(fitted, residuals, alpha=0.6, edgecolors='black', linewidths=0.5) axes[0, 0].axhline(y=0, color='r', linestyle='--') axes[0, 0].set_xlabel('Fitted values') axes[0, 0].set_ylabel('Residuals') axes[0, 0].set_title('Residuals vs Fitted') axes[0, 0].grid(alpha=0.3) stats.probplot(residuals, dist="norm", plot=axes[0, 1]) axes[0, 1].set_title('Normal Q-Q') axes[0, 1].grid(alpha=0.3) std_resid = residuals / residuals.std() axes[1, 0].scatter(fitted, np.sqrt(np.abs(std_resid)), alpha=0.6, edgecolors='black', linewidths=0.5) axes[1, 0].set_xlabel('Fitted values') axes[1, 0].set_ylabel('√|Standardized residuals|') axes[1, 0].set_title('Scale-Location') axes[1, 0].grid(alpha=0.3) axes[1, 1].hist(residuals, bins='auto', edgecolor='black', alpha=0.7) axes[1, 1].set_xlabel('Residuals') axes[1, 1].set_ylabel('Frequency') axes[1, 1].set_title('Histogram of Residuals') axes[1, 1].grid(alpha=0.3, axis='y') plt.tight_layout() plt.show() issues = [] if sw_p <= alpha: issues.append("non-normal residuals (consider transformation or robust/bootstrap inference)") if bp_p <= alpha: issues.append("heteroscedasticity (consider robust SEs, e.g. model.get_robustcov_results('HC3'), or WLS)") if dw < 1.5 or dw > 2.5: issues.append("possible autocorrelation (if data are ordered/temporal, use time-series methods or cluster-robust SEs)") if not np.isnan(max_vif) and max_vif > 5: issues.append("multicollinearity (VIF > 5; consider dropping/combining predictors)") return { 'residual_normality': {'test': 'Shapiro-Wilk', 'statistic': sw_stat, 'p_value': sw_p, 'ok': sw_p > alpha}, 'heteroscedasticity': {'test': 'Breusch-Pagan', 'statistic': bp_lm, 'p_value': bp_p, 'ok': bp_p > alpha}, 'autocorrelation': {'test': 'Durbin-Watson', 'statistic': dw, 'ok': 1.5 <= dw <= 2.5}, 'vif': vif_table, 'max_vif': max_vif, 'issues': issues, 'interpretation': ("No assumption concerns detected." if not issues else "Concerns: " + "; ".join(issues)) } def comprehensive_assumption_check( data: pd.DataFrame, value_col: str, group_col: Optional[str] = None, alpha: float = 0.05 ) -> Dict: """ Perform comprehensive assumption checking for common statistical tests. Parameters ---------- data : pd.DataFrame Data to check value_col : str Column name for dependent variable group_col : str, optional Column name for grouping variable (if applicable) alpha : float Significance level Returns ------- dict Summary of all assumption checks """ print("=" * 70) print("COMPREHENSIVE ASSUMPTION CHECK") print("=" * 70) results = {} # Outlier detection print("\n1. OUTLIER DETECTION") print("-" * 70) outlier_results = detect_outliers( data[value_col].dropna(), name=value_col, method='iqr', plot=True ) results['outliers'] = outlier_results print(f" {outlier_results['interpretation']}") print(f" {outlier_results['recommendation']}") # Check if grouped data if group_col is not None: # Normality per group print(f"\n2. NORMALITY CHECK (by {group_col})") print("-" * 70) normality_results = check_normality_per_group( data, value_col, group_col, alpha=alpha, plot=True ) results['normality_per_group'] = normality_results print(normality_results.to_string(index=False)) all_normal = normality_results['Normal'].eq('Yes').all() print(f"\n All groups normal: {'Yes' if all_normal else 'No'}") if not all_normal: print(" → Consider non-parametric alternative (Mann-Whitney, Kruskal-Wallis)") # Homogeneity of variance print(f"\n3. HOMOGENEITY OF VARIANCE") print("-" * 70) homogeneity_results = check_homogeneity_of_variance( data, value_col, group_col, alpha=alpha, plot=True ) results['homogeneity'] = homogeneity_results print(f" {homogeneity_results['interpretation']}") print(f" {homogeneity_results['recommendation']}") else: # Overall normality print(f"\n2. NORMALITY CHECK") print("-" * 70) normality_results = check_normality( data[value_col].dropna(), name=value_col, alpha=alpha, plot=True ) results['normality'] = normality_results print(f" {normality_results['interpretation']}") print(f" {normality_results['recommendation']}") # Summary print("\n" + "=" * 70) print("SUMMARY") print("=" * 70) if group_col is not None: all_normal = results.get('normality_per_group', pd.DataFrame()).get('Normal', pd.Series()).eq('Yes').all() is_homogeneous = results.get('homogeneity', {}).get('is_homogeneous', False) if all_normal and is_homogeneous: print("✓ All assumptions met. Proceed with parametric test (t-test, ANOVA).") elif not all_normal: print("✗ Normality violated. Use non-parametric alternative.") elif not is_homogeneous: print("✗ Homogeneity violated. Use Welch's correction or transformation.") else: is_normal = results.get('normality', {}).get('is_normal', False) if is_normal: print("✓ Normality assumption met.") else: print("✗ Normality violated. Consider transformation or non-parametric method.") print("=" * 70) return results if __name__ == "__main__": # Example usage np.random.seed(42) # Simulate data group_a = np.random.normal(75, 8, 50) group_b = np.random.normal(68, 10, 50) df = pd.DataFrame({ 'score': np.concatenate([group_a, group_b]), 'group': ['A'] * 50 + ['B'] * 50 }) # Run comprehensive check results = comprehensive_assumption_check( df, value_col='score', group_col='group', alpha=0.05 )