"""PyMC Model Diagnostics Script Comprehensive diagnostic checks for PyMC models.Run this after sampling to validate results before interpretation. Usage: from scripts.model_diagnostics import check_diagnostics, create_diagnostic_report # Quick check check_diagnostics(idata) # Full report with plots create_diagnostic_report(idata, var_names=['alpha', 'beta', 'sigma'], output_dir='diagnostics/')""" import arviz as azimport numpy as npimport matplotlib.pyplot as pltfrom pathlib import Path def check_diagnostics(idata, var_names=None, ess_threshold=400, rhat_threshold=1.01): """ Perform comprehensive diagnostic checks on MCMC samples. Parameters ---------- idata : arviz.InferenceData InferenceData object from pm.sample() var_names : list, optional Variables to check. If None, checks all model parameters ess_threshold : int Minimum acceptable effective sample size (default: 400) rhat_threshold : float Maximum acceptable R-hat value (default: 1.01) Returns ------- dict Dictionary with diagnostic results and flags """ print("="*70) print(" " * 20 + "MCMC DIAGNOSTICS REPORT") print("="*70) # Get summary statistics summary = az.summary(idata, var_names=var_names) results = { 'summary': summary, 'has_issues': False, 'issues': [] } # 1. Check R-hat (convergence) print("\n1. CONVERGENCE CHECK (R-hat)") print("-" * 70) bad_rhat = summary[summary['r_hat'] > rhat_threshold] if len(bad_rhat) > 0: print(f"⚠️ WARNING: {len(bad_rhat)} parameters have R-hat > {rhat_threshold}") print("\nTop 10 worst R-hat values:") print(bad_rhat[['r_hat']].sort_values('r_hat', ascending=False).head(10)) print("\n⚠️ Chains may not have converged!") print(" → Run longer chains or check for multimodality") results['has_issues'] = True results['issues'].append('convergence') else: print(f"✓ All R-hat values ≤ {rhat_threshold}") print(" Chains have converged successfully") # 2. Check Effective Sample Size print("\n2. EFFECTIVE SAMPLE SIZE (ESS)") print("-" * 70) low_ess_bulk = summary[summary['ess_bulk'] < ess_threshold] low_ess_tail = summary[summary['ess_tail'] < ess_threshold] if len(low_ess_bulk) > 0 or len(low_ess_tail) > 0: print(f"⚠️ WARNING: Some parameters have ESS < {ess_threshold}") if len(low_ess_bulk) > 0: print(f"\n Bulk ESS issues ({len(low_ess_bulk)} parameters):") print(low_ess_bulk[['ess_bulk']].sort_values('ess_bulk').head(10)) if len(low_ess_tail) > 0: print(f"\n Tail ESS issues ({len(low_ess_tail)} parameters):") print(low_ess_tail[['ess_tail']].sort_values('ess_tail').head(10)) print("\n⚠️ High autocorrelation detected!") print(" → Sample more draws or reparameterize to reduce correlation") results['has_issues'] = True results['issues'].append('low_ess') else: print(f"✓ All ESS values ≥ {ess_threshold}") print(" Sufficient effective samples") # 3. Check Divergences print("\n3. DIVERGENT TRANSITIONS") print("-" * 70) divergences = idata.sample_stats.diverging.sum().item() if divergences > 0: total_samples = len(idata.posterior.draw) * len(idata.posterior.chain) divergence_rate = divergences / total_samples * 100 print(f"⚠️ WARNING: {divergences} divergent transitions ({divergence_rate:.2f}% of samples)") print("\n Divergences indicate biased sampling in difficult posterior regions") print(" Solutions:") print(" → Increase target_accept (e.g., target_accept=0.95 or 0.99)") print(" → Use non-centered parameterization for hierarchical models") print(" → Add stronger/more informative priors") print(" → Check for model misspecification") results['has_issues'] = True results['issues'].append('divergences') results['n_divergences'] = divergences else: print("✓ No divergences detected") print(" NUTS explored the posterior successfully") # 4. Check Tree Depth print("\n4. TREE DEPTH") print("-" * 70) tree_depth = idata.sample_stats.tree_depth max_tree_depth = tree_depth.max().item() # Typical max_treedepth is 10 (default in PyMC) hits_max = (tree_depth >= 10).sum().item() if hits_max > 0: total_samples = len(idata.posterior.draw) * len(idata.posterior.chain) hit_rate = hits_max / total_samples * 100 print(f"⚠️ WARNING: Hit maximum tree depth {hits_max} times ({hit_rate:.2f}% of samples)") print("\n Model may be difficult to explore efficiently") print(" Solutions:") print(" → Reparameterize model to improve geometry") print(" → Increase max_treedepth (if necessary)") results['issues'].append('max_treedepth') else: print(f"✓ No maximum tree depth issues") print(f" Maximum tree depth reached: {max_tree_depth}") # 5. Check Energy (if available) if hasattr(idata.sample_stats, 'energy'): print("\n5. ENERGY DIAGNOSTICS") print("-" * 70) print("✓ Energy statistics available") print(" Use az.plot_energy(idata) to visualize energy transitions") print(" Good separation indicates healthy HMC sampling") # Summary print("\n" + "="*70) print("SUMMARY") print("="*70) if not results['has_issues']: print("✓ All diagnostics passed!") print(" Your model has sampled successfully.") print(" Proceed with inference and interpretation.") else: print("⚠️ Some diagnostics failed!") print(f" Issues found: {', '.join(results['issues'])}") print(" Review warnings above and consider re-running with adjustments.") print("="*70) return results def create_diagnostic_report(idata, var_names=None, output_dir='diagnostics/', show=False): """ Create comprehensive diagnostic report with plots. Parameters ---------- idata : arviz.InferenceData InferenceData object from pm.sample() var_names : list, optional Variables to plot. If None, uses all model parameters output_dir : str Directory to save diagnostic plots show : bool Whether to display plots (default: False, just save) Returns ------- dict Diagnostic results from check_diagnostics """ # Create output directory output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) # Run diagnostic checks results = check_diagnostics(idata, var_names=var_names) print(f"\nGenerating diagnostic plots in '{output_dir}'...") # 1. Trace plots fig, axes = plt.subplots( len(var_names) if var_names else 5, 2, figsize=(12, 10) ) az.plot_trace(idata, var_names=var_names, axes=axes) plt.tight_layout() plt.savefig(output_path / 'trace_plots.png', dpi=300, bbox_inches='tight') print(f" ✓ Saved trace plots") if show: plt.show() else: plt.close() # 2. Rank plots (check mixing) fig = plt.figure(figsize=(12, 8)) az.plot_rank(idata, var_names=var_names) plt.tight_layout() plt.savefig(output_path / 'rank_plots.png', dpi=300, bbox_inches='tight') print(f" ✓ Saved rank plots") if show: plt.show() else: plt.close() # 3. Autocorrelation plots fig = plt.figure(figsize=(12, 8)) az.plot_autocorr(idata, var_names=var_names, combined=True) plt.tight_layout() plt.savefig(output_path / 'autocorr_plots.png', dpi=300, bbox_inches='tight') print(f" ✓ Saved autocorrelation plots") if show: plt.show() else: plt.close() # 4. Energy plot (if available) if hasattr(idata.sample_stats, 'energy'): fig = plt.figure(figsize=(10, 6)) az.plot_energy(idata) plt.tight_layout() plt.savefig(output_path / 'energy_plot.png', dpi=300, bbox_inches='tight') print(f" ✓ Saved energy plot") if show: plt.show() else: plt.close() # 5. ESS plot fig = plt.figure(figsize=(10, 6)) az.plot_ess(idata, var_names=var_names, kind='evolution') plt.tight_layout() plt.savefig(output_path / 'ess_evolution.png', dpi=300, bbox_inches='tight') print(f" ✓ Saved ESS evolution plot") if show: plt.show() else: plt.close() # Save summary to CSV results['summary'].to_csv(output_path / 'summary_statistics.csv') print(f" ✓ Saved summary statistics") print(f"\nDiagnostic report complete! Files saved in '{output_dir}'") return results def compare_prior_posterior(idata, prior_idata, var_names=None, output_path=None): """ Compare prior and posterior distributions. Parameters ---------- idata : arviz.InferenceData InferenceData with posterior samples prior_idata : arviz.InferenceData InferenceData with prior samples var_names : list, optional Variables to compare output_path : str, optional If provided, save plot to this path Returns ------- None """ fig, axes = plt.subplots( len(var_names) if var_names else 3, 1, figsize=(10, 8) ) if not isinstance(axes, np.ndarray): axes = [axes] for idx, var in enumerate(var_names if var_names else list(idata.posterior.data_vars)[:3]): # Plot prior az.plot_dist( prior_idata.prior[var].values.flatten(), label='Prior', ax=axes[idx], color='blue', alpha=0.3 ) # Plot posterior az.plot_dist( idata.posterior[var].values.flatten(), label='Posterior', ax=axes[idx], color='green', alpha=0.3 ) axes[idx].set_title(f'{var}: Prior vs Posterior') axes[idx].legend() plt.tight_layout() if output_path: plt.savefig(output_path, dpi=300, bbox_inches='tight') print(f"Prior-posterior comparison saved to {output_path}") else: plt.show() # Example usageif __name__ == '__main__': print("This script provides diagnostic functions for PyMC models.") print("\nExample usage:") print(""" import pymc as pm from scripts.model_diagnostics import check_diagnostics, create_diagnostic_report # After sampling with pm.Model() as model: # ... define model ... idata = pm.sample() # Quick diagnostic check results = check_diagnostics(idata) # Full diagnostic report with plots create_diagnostic_report( idata, var_names=['alpha', 'beta', 'sigma'], output_dir='my_diagnostics/' ) """)
"""
PyMC Model Diagnostics Script
Comprehensive diagnostic checks for PyMC models.
Run this after sampling to validate results before interpretation.
Usage:
from scripts.model_diagnostics import check_diagnostics, create_diagnostic_report
# Quick check
check_diagnostics(idata)
# Full report with plots
create_diagnostic_report(idata, var_names=['alpha', 'beta', 'sigma'], output_dir='diagnostics/')
"""
import arviz as az
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
def check_diagnostics(idata, var_names=None, ess_threshold=400, rhat_threshold=1.01):
"""
Perform comprehensive diagnostic checks on MCMC samples.
Parameters
----------
idata : arviz.InferenceData
InferenceData object from pm.sample()
var_names : list, optional
Variables to check. If None, checks all model parameters
ess_threshold : int
Minimum acceptable effective sample size (default: 400)
rhat_threshold : float
Maximum acceptable R-hat value (default: 1.01)
Returns
-------
dict
Dictionary with diagnostic results and flags
"""
print("="*70)
print(" " * 20 + "MCMC DIAGNOSTICS REPORT")
print("="*70)
# Get summary statistics
summary = az.summary(idata, var_names=var_names)
results = {
'summary': summary,
'has_issues': False,
'issues': []
}
# 1. Check R-hat (convergence)
print("\n1. CONVERGENCE CHECK (R-hat)")
print("-" * 70)
bad_rhat = summary[summary['r_hat'] > rhat_threshold]
if len(bad_rhat) > 0:
print(f"⚠️ WARNING: {len(bad_rhat)} parameters have R-hat > {rhat_threshold}")
print("\nTop 10 worst R-hat values:")
print(bad_rhat[['r_hat']].sort_values('r_hat', ascending=False).head(10))
print("\n⚠️ Chains may not have converged!")
print(" → Run longer chains or check for multimodality")
results['has_issues'] = True
results['issues'].append('convergence')
else:
print(f"✓ All R-hat values ≤ {rhat_threshold}")
print(" Chains have converged successfully")
# 2. Check Effective Sample Size
print("\n2. EFFECTIVE SAMPLE SIZE (ESS)")
print("-" * 70)
low_ess_bulk = summary[summary['ess_bulk'] < ess_threshold]
low_ess_tail = summary[summary['ess_tail'] < ess_threshold]
if len(low_ess_bulk) > 0 or len(low_ess_tail) > 0:
print(f"⚠️ WARNING: Some parameters have ESS < {ess_threshold}")
if len(low_ess_bulk) > 0:
print(f"\n Bulk ESS issues ({len(low_ess_bulk)} parameters):")
print(low_ess_bulk[['ess_bulk']].sort_values('ess_bulk').head(10))
if len(low_ess_tail) > 0:
print(f"\n Tail ESS issues ({len(low_ess_tail)} parameters):")
print(low_ess_tail[['ess_tail']].sort_values('ess_tail').head(10))
print("\n⚠️ High autocorrelation detected!")
print(" → Sample more draws or reparameterize to reduce correlation")
results['has_issues'] = True
results['issues'].append('low_ess')
else:
print(f"✓ All ESS values ≥ {ess_threshold}")
print(" Sufficient effective samples")
# 3. Check Divergences
print("\n3. DIVERGENT TRANSITIONS")
print("-" * 70)
divergences = idata.sample_stats.diverging.sum().item()
if divergences > 0:
total_samples = len(idata.posterior.draw) * len(idata.posterior.chain)
divergence_rate = divergences / total_samples * 100
print(f"⚠️ WARNING: {divergences} divergent transitions ({divergence_rate:.2f}% of samples)")
print("\n Divergences indicate biased sampling in difficult posterior regions")
print(" Solutions:")
print(" → Increase target_accept (e.g., target_accept=0.95 or 0.99)")
print(" → Use non-centered parameterization for hierarchical models")
print(" → Add stronger/more informative priors")
print(" → Check for model misspecification")
results['has_issues'] = True
results['issues'].append('divergences')
results['n_divergences'] = divergences
else:
print("✓ No divergences detected")
print(" NUTS explored the posterior successfully")
# 4. Check Tree Depth
print("\n4. TREE DEPTH")
print("-" * 70)
tree_depth = idata.sample_stats.tree_depth
max_tree_depth = tree_depth.max().item()
# Typical max_treedepth is 10 (default in PyMC)
hits_max = (tree_depth >= 10).sum().item()
if hits_max > 0:
total_samples = len(idata.posterior.draw) * len(idata.posterior.chain)
hit_rate = hits_max / total_samples * 100
print(f"⚠️ WARNING: Hit maximum tree depth {hits_max} times ({hit_rate:.2f}% of samples)")
print("\n Model may be difficult to explore efficiently")
print(" Solutions:")
print(" → Reparameterize model to improve geometry")
print(" → Increase max_treedepth (if necessary)")
results['issues'].append('max_treedepth')
else:
print(f"✓ No maximum tree depth issues")
print(f" Maximum tree depth reached: {max_tree_depth}")
# 5. Check Energy (if available)
if hasattr(idata.sample_stats, 'energy'):
print("\n5. ENERGY DIAGNOSTICS")
print("-" * 70)
print("✓ Energy statistics available")
print(" Use az.plot_energy(idata) to visualize energy transitions")
print(" Good separation indicates healthy HMC sampling")
# Summary
print("\n" + "="*70)
print("SUMMARY")
print("="*70)
if not results['has_issues']:
print("✓ All diagnostics passed!")
print(" Your model has sampled successfully.")
print(" Proceed with inference and interpretation.")
else:
print("⚠️ Some diagnostics failed!")
print(f" Issues found: {', '.join(results['issues'])}")
print(" Review warnings above and consider re-running with adjustments.")
print("="*70)
return results
def create_diagnostic_report(idata, var_names=None, output_dir='diagnostics/', show=False):
"""
Create comprehensive diagnostic report with plots.
Parameters
----------
idata : arviz.InferenceData
InferenceData object from pm.sample()
var_names : list, optional
Variables to plot. If None, uses all model parameters
output_dir : str
Directory to save diagnostic plots
show : bool
Whether to display plots (default: False, just save)
Returns
-------
dict
Diagnostic results from check_diagnostics
"""
# Create output directory
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Run diagnostic checks
results = check_diagnostics(idata, var_names=var_names)
print(f"\nGenerating diagnostic plots in '{output_dir}'...")
# 1. Trace plots
fig, axes = plt.subplots(
len(var_names) if var_names else 5,
2,
figsize=(12, 10)
)
az.plot_trace(idata, var_names=var_names, axes=axes)
plt.tight_layout()
plt.savefig(output_path / 'trace_plots.png', dpi=300, bbox_inches='tight')
print(f" ✓ Saved trace plots")
if show:
plt.show()
else:
plt.close()
# 2. Rank plots (check mixing)
fig = plt.figure(figsize=(12, 8))
az.plot_rank(idata, var_names=var_names)
plt.tight_layout()
plt.savefig(output_path / 'rank_plots.png', dpi=300, bbox_inches='tight')
print(f" ✓ Saved rank plots")
if show:
plt.show()
else:
plt.close()
# 3. Autocorrelation plots
fig = plt.figure(figsize=(12, 8))
az.plot_autocorr(idata, var_names=var_names, combined=True)
plt.tight_layout()
plt.savefig(output_path / 'autocorr_plots.png', dpi=300, bbox_inches='tight')
print(f" ✓ Saved autocorrelation plots")
if show:
plt.show()
else:
plt.close()
# 4. Energy plot (if available)
if hasattr(idata.sample_stats, 'energy'):
fig = plt.figure(figsize=(10, 6))
az.plot_energy(idata)
plt.tight_layout()
plt.savefig(output_path / 'energy_plot.png', dpi=300, bbox_inches='tight')
print(f" ✓ Saved energy plot")
if show:
plt.show()
else:
plt.close()
# 5. ESS plot
fig = plt.figure(figsize=(10, 6))
az.plot_ess(idata, var_names=var_names, kind='evolution')
plt.tight_layout()
plt.savefig(output_path / 'ess_evolution.png', dpi=300, bbox_inches='tight')
print(f" ✓ Saved ESS evolution plot")
if show:
plt.show()
else:
plt.close()
# Save summary to CSV
results['summary'].to_csv(output_path / 'summary_statistics.csv')
print(f" ✓ Saved summary statistics")
print(f"\nDiagnostic report complete! Files saved in '{output_dir}'")
return results
def compare_prior_posterior(idata, prior_idata, var_names=None, output_path=None):
"""
Compare prior and posterior distributions.
Parameters
----------
idata : arviz.InferenceData
InferenceData with posterior samples
prior_idata : arviz.InferenceData
InferenceData with prior samples
var_names : list, optional
Variables to compare
output_path : str, optional
If provided, save plot to this path
Returns
-------
None
"""
fig, axes = plt.subplots(
len(var_names) if var_names else 3,
1,
figsize=(10, 8)
)
if not isinstance(axes, np.ndarray):
axes = [axes]
for idx, var in enumerate(var_names if var_names else list(idata.posterior.data_vars)[:3]):
# Plot prior
az.plot_dist(
prior_idata.prior[var].values.flatten(),
label='Prior',
ax=axes[idx],
color='blue',
alpha=0.3
)
# Plot posterior
az.plot_dist(
idata.posterior[var].values.flatten(),
label='Posterior',
ax=axes[idx],
color='green',
alpha=0.3
)
axes[idx].set_title(f'{var}: Prior vs Posterior')
axes[idx].legend()
plt.tight_layout()
if output_path:
plt.savefig(output_path, dpi=300, bbox_inches='tight')
print(f"Prior-posterior comparison saved to {output_path}")
else:
plt.show()
# Example usage
if __name__ == '__main__':
print("This script provides diagnostic functions for PyMC models.")
print("\nExample usage:")
print("""
import pymc as pm
from scripts.model_diagnostics import check_diagnostics, create_diagnostic_report
# After sampling
with pm.Model() as model:
# ... define model ...
idata = pm.sample()
# Quick diagnostic check
results = check_diagnostics(idata)
# Full diagnostic report with plots
create_diagnostic_report(
idata,
var_names=['alpha', 'beta', 'sigma'],
output_dir='my_diagnostics/'
)
""")