Skills › Scientific Agent Skills › pymc › model_diagnostics.py
model_diagnostics.py
262 lines · skills/pymc/scripts/model_diagnostics.py at 5f2bc03
1"""
2PyMC Model Diagnostics Script
4Comprehensive diagnostic checks for PyMC models.
5Run this after sampling to validate results before interpretation.
7Usage:
8 from scripts.model_diagnostics import check_diagnostics, create_diagnostic_report
10 # Quick check
11 check_diagnostics(idata)
13 # Full report with plots
14 create_diagnostic_report(idata, var_names=['alpha', 'beta', 'sigma'], output_dir='diagnostics/')
15"""
17import arviz as az
18import matplotlib.pyplot as plt
19from pathlib import Path
22def check_diagnostics(idata, var_names=None, ess_threshold=400, rhat_threshold=1.01,
23 max_treedepth=None, bfmi_threshold=0.3):
24 """Screen MCMC draws; a clean screen is not proof of convergence or identification.
26 Targets PyMC 6 / ArviZ 1 DataTree objects. ESS is pooled across chains.
27 Supply the sampler's configured max_treedepth when it does not record
28 reached_max_treedepth. Missing HMC statistics are reported as unavailable,
29 not silently counted as zero. VI draws are not independent MCMC chains.
30 """
31 import numpy as np
33 if ess_threshold <= 0 or not np.isfinite(ess_threshold):
34 raise ValueError("ess_threshold must be finite and positive")
35 if not np.isfinite(rhat_threshold) or rhat_threshold < 1:
36 raise ValueError("rhat_threshold must be finite and at least 1")
37 if max_treedepth is not None and (max_treedepth <= 0 or int(max_treedepth) != max_treedepth):
38 raise ValueError("max_treedepth must be a positive integer or None")
39 if not np.isfinite(bfmi_threshold) or bfmi_threshold <= 0:
40 raise ValueError("bfmi_threshold must be finite and positive")
41 if not hasattr(idata, 'posterior') or not idata.posterior.data_vars:
42 raise ValueError("A nonempty posterior group is required")
43 summary = az.summary(idata, var_names=var_names, round_to="none")
44 if summary.empty:
45 raise ValueError("No posterior variables selected")
46 issues, unavailable = [], []
47 results = {'summary': summary, 'issues': issues, 'unavailable': unavailable}
48 if idata.posterior.sizes.get('chain', 0) < 2:
49 issues.append('insufficient_chains')
50 diagnostic_columns = ['r_hat', 'ess_bulk', 'ess_tail']
51 if not np.isfinite(summary[diagnostic_columns].to_numpy()).all():
52 issues.append('nonfinite_diagnostics')
53 if (summary['r_hat'] > rhat_threshold).any():
54 issues.append('convergence')
55 if (summary[['ess_bulk', 'ess_tail']] < ess_threshold).any().any():
56 issues.append('low_ess')
58 stats = getattr(idata, 'sample_stats', None)
59 if stats is not None and 'diverging' in stats:
60 divergences = stats['diverging'].values
61 if not np.isfinite(divergences).all():
62 issues.append('nonfinite_sampler_stats')
63 else:
64 count = int(divergences.sum())
65 if count:
66 results['n_divergences'] = count
67 issues.append('divergences')
68 else:
69 unavailable.append('divergences')
71 if stats is not None and 'reached_max_treedepth' in stats:
72 hit_values = stats['reached_max_treedepth'].values
73 if not np.isfinite(hit_values).all():
74 issues.append('nonfinite_sampler_stats')
75 elif hit_values.any():
76 issues.append('max_treedepth')
77 elif stats is not None and 'tree_depth' in stats and max_treedepth is not None:
78 depth = stats['tree_depth'].values
79 if not np.isfinite(depth).all():
80 issues.append('nonfinite_sampler_stats')
81 elif (depth >= max_treedepth).any():
82 issues.append('max_treedepth')
83 else:
84 unavailable.append('max_treedepth')
86 if stats is not None and 'energy' in stats:
87 # ArviZ 1.x returns a DataTree with one energy BFMI value per chain.
88 bfmi = np.asarray(az.bfmi(idata)['energy'].values)
89 results['bfmi'] = bfmi
90 if not np.isfinite(bfmi).all():
91 issues.append('nonfinite_bfmi')
92 elif (bfmi < bfmi_threshold).any():
93 issues.append('low_bfmi')
94 else:
95 unavailable.append('bfmi')
96 results['has_issues'] = bool(issues)
97 print('[WARN] Diagnostic flags: ' + ', '.join(issues) if issues else
98 '[OK] No flags in the available numerical diagnostics')
99 if unavailable:
100 print('[INFO] Unavailable checks: ' + ', '.join(unavailable))
101 print('Inspect traces, MCSE for the estimand, prior sensitivity and predictive checks.')
102 print('This screen does not establish convergence, identifiability or model validity.')
103 return results
106def create_diagnostic_report(idata, var_names=None, output_dir='diagnostics/', show=False):
107 """
108 Create comprehensive diagnostic report with plots.
110 Parameters
111 ----------
112 idata : xarray.DataTree
113 Posterior object from pm.sample()
114 var_names : list, optional
115 Variables to plot. If None, uses all model parameters
116 output_dir : str
117 Directory to save diagnostic plots
118 show : bool
119 Whether to display plots (default: False, just save)
121 Returns
122 -------
123 dict
124 Diagnostic results from check_diagnostics
125 """
126 # Create output directory
127 output_path = Path(output_dir)
128 output_path.mkdir(parents=True, exist_ok=True)
130 # Run diagnostic checks
131 results = check_diagnostics(idata, var_names=var_names)
133 print(f"\nGenerating diagnostic plots in '{output_dir}'...")
135 # ArviZ 1.x plots return a PlotCollection and do not draw into pyplot's
136 # current figure, so the figure must be saved through the collection --
137 # plt.savefig() would write a blank image.
138 def _save(plot_collection, filename, label):
139 plot_collection.savefig(
140 output_path / filename, dpi=300, bbox_inches='tight'
141 )
142 print(f" [OK] Saved {label}")
143 if show:
144 plt.show()
145 else:
146 plt.close(plot_collection.viz['figure'].item())
148 # 1. Trace plots
149 _save(
150 az.plot_trace_dist(idata, var_names=var_names),
151 'trace_plots.png',
152 'trace plots',
153 )
155 # 2. Rank plots (check mixing)
156 _save(
157 az.plot_rank(idata, var_names=var_names),
158 'rank_plots.png',
159 'rank plots',
160 )
162 # 3. Autocorrelation plots
163 _save(
164 az.plot_autocorr(idata, var_names=var_names),
165 'autocorr_plots.png',
166 'autocorrelation plots',
167 )
169 # 4. Energy plot (if available)
170 if hasattr(idata, 'sample_stats') and 'energy' in idata.sample_stats:
171 _save(az.plot_energy(idata), 'energy_plot.png', 'energy plot')
173 # 5. ESS plot. ArviZ 1.x offers 'local' and 'quantile'; the old
174 # 'evolution' kind no longer exists.
175 _save(
176 az.plot_ess(idata, var_names=var_names, kind='local'),
177 'ess_local.png',
178 'local ESS plot',
179 )
181 # Save summary to CSV
182 results['summary'].to_csv(output_path / 'summary_statistics.csv')
183 print(f" [OK] Saved summary statistics")
185 print(f"\nDiagnostic report complete! Files saved in '{output_dir}'")
187 return results
190def compare_prior_posterior(idata, prior_idata, var_names=None, output_path=None):
191 """
192 Compare prior and posterior distributions.
194 Parameters
195 ----------
196 idata : xarray.DataTree
197 Posterior object with posterior samples
198 prior_idata : xarray.DataTree
199 Prior object with prior samples
200 var_names : list, optional
201 Variables to compare. Defaults to the first three posterior variables.
202 output_path : str, optional
203 If provided, save plot to this path
205 Returns
206 -------
207 arviz_plots.PlotCollection
208 The overlaid prior/posterior figure.
209 """
210 if var_names is None:
211 var_names = list(idata.posterior.data_vars)[:3]
213 # ArviZ 1.x plot functions take a DataTree and a group name, not a flat
214 # array plus an axis. Passing the first call's PlotCollection back into the
215 # second is what overlays the two distributions on the same axes.
216 collection = az.plot_dist(
217 prior_idata,
218 group='prior',
219 var_names=var_names,
220 visuals={'dist': {'color': 'blue'}},
221 )
222 az.plot_dist(
223 idata,
224 group='posterior',
225 var_names=var_names,
226 plot_collection=collection,
227 visuals={'dist': {'color': 'green'}},
228 )
230 if output_path:
231 collection.savefig(output_path, dpi=300, bbox_inches='tight')
232 print(f"Prior-posterior comparison saved to {output_path}")
233 print("Prior is blue, posterior is green")
234 else:
235 plt.show()
237 return collection
240# Example usage
241if __name__ == '__main__':
242 print("This script provides diagnostic functions for PyMC models.")
243 print("\nExample usage:")
244 print("""
245 import pymc as pm
246 from scripts.model_diagnostics import check_diagnostics, create_diagnostic_report
248 # After sampling
249 with pm.Model() as model:
250 # ... define model ...
251 idata = pm.sample()
253 # Quick diagnostic check
254 results = check_diagnostics(idata)
256 # Full diagnostic report with plots
257 create_diagnostic_report(
258 idata,
259 var_names=['alpha', 'beta', 'sigma'],
260 output_dir='my_diagnostics/'
261 )
262 """)