Skills › Scientific Agent Skills › pymc › model_comparison.py
model_comparison.py
345 lines · skills/pymc/scripts/model_comparison.py at 5f2bc03
1"""
2PyMC Model Comparison Script
4Utilities for comparing multiple Bayesian models using information criteria
5and cross-validation metrics.
7Usage:
8 from scripts.model_comparison import compare_models, plot_model_comparison
10 # Compare multiple models
11 comparison = compare_models(
12 {'model1': idata1, 'model2': idata2, 'model3': idata3},
13 ic='loo'
14 )
16 # Visualize comparison
17 plot_model_comparison(comparison, output_path='model_comparison.png')
18"""
20import arviz as az
21import numpy as np
22import pandas as pd
23import matplotlib.pyplot as plt
24from typing import Any, Dict
27#: ArviZ 1.x compares models on PSIS-LOO ELPD only; there is no `ic=` switch and
28#: no deviance scale. WAIC is still available on its own via `az.waic()`.
29SUPPORTED_IC = ('loo', 'elpd')
32def compare_models(models_dict: Dict[str, Any],
33 ic='loo',
34 verbose=True, var_name=None):
35 """
36 Compare multiple models by expected log pointwise predictive density.
38 Parameters
39 ----------
40 models_dict : dict
41 Dictionary mapping model names to PyMC posterior objects.
42 All models must have log_likelihood computed.
43 ic : str
44 Information criterion. Only 'loo' (equivalently 'elpd') is supported:
45 ArviZ 1.x ranks models on PSIS-LOO ELPD.
46 verbose : bool
47 Print detailed comparison results (default: True)
49 Returns
50 -------
51 pd.DataFrame
52 Comparison DataFrame with model rankings and statistics, on the ELPD
53 scale (higher is better, so `elpd_diff` is 0 for the best model and
54 negative for the others).
56 Notes
57 -----
58 Models must have a log_likelihood group, computed during sampling or
59 afterwards with pm.compute_log_likelihood(idata).
60 """
61 if ic.lower() not in SUPPORTED_IC:
62 raise ValueError(
63 f"unknown information criterion {ic!r}: ArviZ 1.x ranks models on "
64 "PSIS-LOO ELPD, so pass ic='loo'. For WAIC, call az.waic() per "
65 "model directly."
66 )
68 if verbose:
69 print("="*70)
70 print(" " * 25 + "MODEL COMPARISON (LOO)")
71 print("="*70)
73 _validate_comparable_observations(models_dict, var_name)
75 # round_to='none' keeps the columns numeric; the default formats them for
76 # display, which turns every comparison below into a string comparison.
77 comparison = az.compare(models_dict, var_name=var_name, round_to='none')
79 if verbose:
80 print("\nModel Rankings:")
81 print("-"*70)
82 print(comparison.to_string())
84 print("\n" + "="*70)
85 print("INTERPRETATION GUIDE")
86 print("="*70)
87 print("- rank: Model ranking (0 = best)")
88 print("- elpd: PSIS-LOO ELPD estimate (higher is better)")
89 print("- p: Effective number of parameters")
90 print("- elpd_diff: ELPD minus the best model's ELPD (0 for the best)")
91 print("- weight: Predictive stacking weight (not model probability)")
92 print("- se: Standard error of the ELPD estimate")
93 print("- dse: Standard error of the difference")
94 print("- p_worse: Probability the model is worse than the best one")
95 print("- diag_elpd: Reliability diagnostic for the ELPD estimate")
97 print("\n" + "="*70)
98 print("MODEL SELECTION GUIDELINES")
99 print("="*70)
101 best_model = comparison.index[0]
102 print(f"\n[OK] Best model: {best_model}")
104 # Check for a clear winner. Vehtari et al. recommend treating an ELPD
105 # difference below 4 as small, and otherwise judging it against the
106 # standard error of the difference.
107 if len(comparison) > 1:
108 delta = abs(comparison.iloc[1]['elpd_diff'])
109 delta_se = comparison.iloc[1]['dse']
111 if delta < 4:
112 print(f" -> Models are SIMILAR (ELPD difference {delta:.1f} < 4)")
113 print(" Consider model averaging or choose based on simplicity")
114 elif delta > 2 * delta_se:
115 print(
116 f" -> Larger estimated predictive score for {best_model} "
117 f"(ELPD difference {delta:.1f} > 2 SE)"
118 )
119 else:
120 print(
121 f" -> Uncertain estimated advantage for {best_model} "
122 f"(ELPD difference {delta:.1f}, within 2 SE)"
123 )
125 # Reliability. ArviZ 1.x reports this per row as a diagnostic string
126 # instead of the old boolean `warning` column.
127 flagged = [
128 name
129 for name, diagnostic in comparison['diag_elpd'].items()
130 if isinstance(diagnostic, str) and diagnostic.strip() not in ('', 'ok')
131 ]
132 if flagged:
133 print("\n[WARN] WARNING: Some models have reliability issues")
134 print(f" Models with warnings: {', '.join(flagged)}")
135 print(" -> Check Pareto-k diagnostics with check_loo_reliability()")
137 return comparison
140def _validate_comparable_observations(models_dict, var_name=None):
141 """Reject mismatched observed values/order; predictive-unit choice remains scientific."""
142 if not models_dict:
143 raise ValueError("At least one model is required")
144 reference = None
145 reference_layout = None
146 selected_name = var_name
147 for name, idata in models_dict.items():
148 if not hasattr(idata, 'log_likelihood') or not hasattr(idata, 'observed_data'):
149 raise ValueError(f"{name}: log_likelihood and observed_data groups are required")
150 names = list(idata.log_likelihood.data_vars)
151 if selected_name is None:
152 if len(names) != 1:
153 raise ValueError("Specify var_name when there are multiple likelihoods")
154 selected_name = names[0]
155 if selected_name not in idata.log_likelihood or selected_name not in idata.observed_data:
156 raise ValueError(f"{name}: missing likelihood/observations for {selected_name}")
157 observed = idata.observed_data[selected_name]
158 loglik = idata.log_likelihood[selected_name]
159 if not np.isfinite(loglik.values).all():
160 raise ValueError(f"{name}: log likelihood contains nonfinite values")
161 if not {'chain', 'draw'}.issubset(loglik.dims):
162 raise ValueError(f"{name}: log likelihood needs chain and draw dimensions")
163 layout = loglik.isel(chain=0, draw=0, drop=True)
164 if reference is None:
165 reference = observed
166 reference_layout = layout
167 elif not observed.identical(reference):
168 raise ValueError("Models must use identical observed outcomes, dimensions and coordinates")
169 elif layout.dims != reference_layout.dims or layout.sizes != reference_layout.sizes or not layout.coords.equals(reference_layout.coords):
170 raise ValueError("Pointwise likelihood dimensions and coordinates must match")
171 return selected_name
174def check_loo_reliability(models_dict: Dict[str, Any], threshold=None,
175 verbose=True, var_name=None):
176 """Report PSIS diagnostics; default to each result's sample-size-aware good_k.
178 A custom threshold may be stricter. Nonfinite k values are always flagged.
179 This checks importance sampling, not convergence or the chosen predictive unit.
180 """
181 if threshold is not None and not np.isfinite(threshold):
182 raise ValueError("threshold must be finite or None")
183 _validate_comparable_observations(models_dict, var_name)
184 results = {}
185 for name, idata in models_dict.items():
186 loo_result = az.loo(idata, pointwise=True, var_name=var_name)
187 pareto_k = np.asarray(loo_result.pareto_k.values)
188 cutoff = float(loo_result.good_k if threshold is None else threshold)
189 invalid = ~np.isfinite(pareto_k)
190 n_high = int(((pareto_k > cutoff) | invalid).sum())
191 results[name] = {
192 'pareto_k': pareto_k, 'n_high': n_high,
193 'n_very_high': int((pareto_k > 1).sum()),
194 'n_nonfinite': int(invalid.sum()), 'threshold': cutoff,
195 'max_k': float(pareto_k.max()), 'loo': loo_result,
196 }
197 if verbose:
198 print(f"{name}: {n_high} observations fail Pareto-k <= {cutoff:.3f}")
199 if n_high:
200 print('[WARN] Investigate influential units; refit LOO cases or use structured K-fold CV.')
201 print('Switching to WAIC does not repair unreliable PSIS-LOO.')
202 return results
205def plot_model_comparison(comparison, output_path=None, show=True):
206 """
207 Visualize model comparison results.
209 Parameters
210 ----------
211 comparison : pd.DataFrame
212 Comparison DataFrame from az.compare()
213 output_path : str, optional
214 If provided, save plot to this path
215 show : bool
216 Whether to display plot (default: True)
218 Returns
219 -------
220 matplotlib.figure.Figure
221 The comparison figure
222 """
223 # ArviZ 1.x returns a PlotCollection and does not draw into pyplot's
224 # current figure, so the figure has to come back out of the collection --
225 # plt.savefig() would write a blank image.
226 collection = az.plot_compare(comparison)
227 fig = collection.viz['figure'].item()
228 fig.suptitle('Model Comparison', fontsize=14, fontweight='bold')
230 if output_path:
231 fig.savefig(output_path, dpi=300, bbox_inches='tight')
232 print(f"Comparison plot saved to {output_path}")
234 if show:
235 plt.show()
236 else:
237 plt.close(fig)
239 return fig
242def model_averaging(models_dict: Dict[str, Any], weights=None, var_name='y_obs',
243 ic='loo', *, group='posterior_predictive', n_samples=None,
244 random_seed=None):
245 """Sample a predictive mixture, preserving within-draw dependence and variance.
247 Returns (draws, weights_in_input_model_order). Draws have shape
248 (n_samples, *prediction_dimensions), not (chain, draw, ...). Each draw selects
249 one model then one complete joint predictive draw. It never averages paired
250 draws. Stacking weights are predictive weights, not posterior model probabilities.
251 All models must provide the same explicit group, variable, dimensions and
252 coordinates; missing predictions and invalid weights raise ValueError.
253 """
254 names = list(models_dict)
255 if not names:
256 raise ValueError("At least one model is required")
257 if weights is None:
258 comparison = compare_models(models_dict, ic=ic, verbose=False)
259 weights = comparison.loc[names, 'weight'].to_numpy()
260 weights = np.asarray(weights, dtype=float)
261 if weights.shape != (len(names),) or not np.isfinite(weights).all() or (weights < 0).any():
262 raise ValueError("Supply one finite nonnegative weight per model")
263 if weights.sum() <= 0 or not np.isfinite(weights.sum()):
264 raise ValueError("Weights must have a finite positive sum")
265 weights = weights / weights.sum()
266 predictions, reference = [], None
267 for name in names:
268 idata = models_dict[name]
269 if not hasattr(idata, group) or var_name not in idata[group]:
270 raise ValueError(f"{name}: missing {group}/{var_name}")
271 pred = idata[group][var_name]
272 if not {'chain', 'draw'}.issubset(pred.dims):
273 raise ValueError(f"{name}: predictions need chain and draw dimensions")
274 dims = [d for d in pred.dims if d not in ('chain', 'draw')]
275 ordered = pred.transpose('chain', 'draw', *dims)
276 layout = ordered.isel(chain=0, draw=0, drop=True)
277 if reference is None:
278 reference = layout
279 elif layout.dims != reference.dims or layout.sizes != reference.sizes or not layout.coords.equals(reference.coords):
280 raise ValueError("Prediction dimensions and coordinates must match exactly")
281 values = ordered.values.reshape((-1, *layout.shape))
282 if not np.isfinite(values).all():
283 raise ValueError(f"{name}: nonfinite predictive draws")
284 predictions.append(values)
285 if n_samples is None:
286 n_samples = min(len(p) for p in predictions)
287 if not isinstance(n_samples, (int, np.integer)) or n_samples <= 0:
288 raise ValueError("n_samples must be a positive integer")
289 rng = np.random.default_rng(random_seed)
290 choices = rng.choice(len(names), size=n_samples, p=weights)
291 mixture = np.empty((n_samples, *reference.shape), dtype=np.result_type(*predictions))
292 for index, pred in enumerate(predictions):
293 mask = choices == index
294 mixture[mask] = pred[rng.integers(len(pred), size=int(mask.sum()))]
295 return mixture, weights
298def cross_validation_comparison(models_dict: Dict[str, Any], k=10, verbose=True):
299 """Print a conceptual K-fold guide; this function does not fit or score models."""
300 if not isinstance(k, int) or k < 2:
301 raise ValueError("k must be an integer >= 2")
302 if verbose:
303 print(f"K-fold guide: {k} folds; no fits have been performed.")
304 print("Choose independent units (rows, groups, or future time blocks) before splitting.")
305 print("For each fold: fit preprocessing and model only on training data; diagnose it.")
306 print("Update predictor AND observed pm.Data containers and coordinates for held-out data.")
307 print("Compute held-out log likelihood with extend_inferencedata=False to preserve training groups.")
308 print("For each held-out unit, compute logsumexp(log_lik over posterior draws) - log(n_draws).")
309 print("Sum these log predictive densities over units; do NOT sum log likelihood over draws.")
310 print("For a joint held-out group, sum log likelihood within that group BEFORE logsumexp.")
311 print("Compare paired fold/unit scores and uncertainty on the same outcome scale.")
314# Example usage
315if __name__ == '__main__':
316 print("This script provides model comparison utilities for PyMC.")
317 print("\nExample usage:")
318 print("""
319 import pymc as pm
320 from scripts.model_comparison import compare_models, check_loo_reliability
322 # Fit multiple models (must include log_likelihood)
323 with pm.Model() as model1:
324 # ... define model 1 ...
325 idata1 = pm.sample()
326 pm.compute_log_likelihood(idata1)
328 with pm.Model() as model2:
329 # ... define model 2 ...
330 idata2 = pm.sample()
331 pm.compute_log_likelihood(idata2)
333 # Compare models
334 models = {'Simple': idata1, 'Complex': idata2}
335 comparison = compare_models(models, ic='loo')
337 # Check reliability
338 reliability = check_loo_reliability(models)
340 # Visualize
341 plot_model_comparison(comparison, output_path='comparison.png')
343 # Model averaging
344 averaged_pred, weights = model_averaging(models, var_name='y_obs')
345 """)