Amortized Bayesian Workflow

Workflow overview

Every amortized Bayesian analysis follows this sequence. Do not skip steps — especially simulator validation and model criticism.

1. Formulate — Define the generative story. What latent variables or parameters generated the observations?
2. Specify the simulator regime — The first iteration always uses offline training for fast turnaround, regardless of simulator speed. The simulator regime only determines the simulation budget for the pilot run:
   • Fast simulator (< 0.05 s per draw): pre-simulate 20 000 datasets, train for 100 epochs
   • Slow simulator (> 1 s – minutes per draw): pre-simulate 3 000–5 000 datasets, train for 100 epochs
   • No simulator / pre-existing bank: use whatever is available; switch to disk training if it does not fit in memory Online training is a refinement step — use it only after the first offline pass shows healthy diagnostics and you want to squeeze out more performance.
3. Define prior + observation model or simulation bank
   • Implement prior and observation model and wrap them in a simulator
   • Pre-simulate the pilot budget into a dict (using workflow.simulate(N)) for offline training
   • If the simulator is external or proprietary, ensure simulations are already generated from the intended prior and data-generating process
4. Choose the architecture — this step is critical; getting it wrong ruins inference. See references/conditioning.md for the full conditioning logic and decision table.
   • "Simple vector" means the observation is a single fixed-length feature vector whose element order is meaningful (e.g., 5 named sensor readings, a pre-computed summary statistic). Only then: route through inference_conditions with no summary network.
   • Set-based / exchangeable data — If the simulator produces N observations that are exchangeable, the data is a set, not a vector. This includes: N i.i.d. draws, regression datasets with (x, y) pairs, repeated measurements, trial-level data, cross-sectional samples. Route through summary_variables with a SetTransformer. Never put this in inference_conditions.
   • Time series — ordered sequences: route through summary_variables with TimeSeriesTransformer or TimeSeriesNetwork.

• Images as conditions / observations for parameter inference — route through summary_variables with ConvolutionalNetwork.
• Images as inferential targets — conditional image generation, spatial field generation, denoising, and other image-valued outputs require an image-capable diffusion inference network. Use bf.networks.DiffusionModel(subnet=...) with UNet, UViT, or ResidualUViT; see references/image-generation.md.
• A workflow can use both slots simultaneously. Fixed-length metadata (e.g., sample size N, scalar design variables) can go in inference_conditions while structured observations go in summary_variables.
• When in doubt, use a summary network. It is always safer to include one than to omit one; a summary network will always be needed if the data has more than one axis.

5. Build the workflow — Prefer bf.BasicWorkflow(...)
   • Decide on which variables to auto-standardize. Prefer standardize="all" unless you have verfied that the simulator outputs are already in a good range for the networks.
6. Run simulation sanity checks — Before training, verify that simulated data look plausible and span the relevant range of real observations. Again, pay attention to what needs to be standardized.
7. Train the amortizer — First iteration always uses offline training for fast feedback:
   • workflow.fit_offline(...) with the pre-simulated pilot budget (default first pass)
   • workflow.fit_online(...) only as a refinement step after offline diagnostics look healthy, or when the user explicitly requests it
   • workflow.fit_disk(...) if streaming simulations from disk Always offer to run training in the terminal so the user can monitor progress interactively.
8. Diagnose in silico — Use held-out simulations with known ground truth using the workflow's built-in diagnostics: workflow.compute_default_diagnostics(...) for numerical results and workflow.plot_default_diagnostics(...) for visual diagnostics.
9. Amortized inference on real data — Use workflow.sample(...)
10. Posterior predictive checks (PPCs) — Re-simulate data from posterior samples and compare to the real data using model-specific test quantities
11. Write a report — Use references/reporting.md to generate a structured report outlining results and next steps.

Hard rules — MUST and NEVER

These rules are non-negotiable. Violating any of them will silently produce wrong results.

• MUST use bf.Adapter() for data routing. Build an explicit adapter chain with .as_set(), .constrain(), .concatenate(), etc. and pass adapter= to BasicWorkflow, as described in references/adapter.md. Do NOT do manual preprocessing — the adapter handles training and inference identically. The naming shorthand (inference_variables=, summary_variables= as kwargs to BasicWorkflow) is ONLY acceptable when the simulator output already has the exact shapes/dtypes the networks expect AND no parameter has bounded support. When in doubt, use an explicit adapter.
• MUST start with the Base network configuration from references/model-sizes.md. Scale up to Large or XL ONLY if diagnostics show poor recovery or calibration after sufficient training. Oversized networks waste compute and can hurt calibration on simple problems.
• MUST use workflow.simulate(N) to generate train/test data — not a Python for-loop over simulator(). The simulator returned by bf.make_simulator is a batched object; workflow.simulate(N) calls it efficiently and returns data in the format the workflow expects.
• MUST use workflow.compute_default_diagnostics(test_data=...) and workflow.plot_default_diagnostics(test_data=...) for in-silico diagnostics. NEVER hand-roll coverage, bias, or calibration computations — the built-in methods are correct, complete, and consistent with the house thresholds.
• For image-valued inference targets, follow references/image-generation.md. MUST use bf.networks.DiffusionModel(subnet=...) with UNet, UViT, or ResidualUViT — not the default low-dimensional setup. Conditions must be spatially concatenable with the image target (broadcast (B, D) to (B, H, W, D)). The standard diagnostic report does not apply; use visual sample grids instead.
• workflow.sample() returns original parameter names, NOT "inference_variables". The adapter's reverse transform restores the original keys from the simulator (e.g., "alpha", "beta", "sigma"). Each parameter has shape (batch, num_samples) for scalars or (batch, num_samples, d) for vectors. NEVER index into "inference_variables" — that key does not exist in the output.
• MUST reuse the existing simulator functions for PPCs. NEVER re-implement the generative model by hand for posterior predictive checks. Loop over a subset of posterior draws (50 is a good default), indexing over the num_samples axis, and pass each draw through the simulator's forward model.
• MUST save history.history as JSON (not CSV, not a DataFrame — it is a plain dict). Then run scripts/inspect_training.py or call inspect_history() in-process.
• MUST pass validation_data= to all fit_* calls. For offline training, hold out ~300 simulations as a separate validation dict. For online training (refinement step only), pass an integer (e.g., validation_data=300) to auto-simulate.
• NEVER mix an explicit adapter= with the naming shorthand (inference_variables=, summary_variables=, inference_conditions= as kwargs). They are mutually exclusive. Passing both causes silent conflicts.
• NEVER flatten structured data into inference_conditions. Sets, time series, and images MUST go through summary_variables with an appropriate summary network.
• workflow.plot_default_diagnostics() ALWAYS returns a dict[str, Figure]. Iterate directly over .items() to save figures. Do not type-check or branch on the return type.
• NEVER skip in-silico diagnostics. Good training loss does not imply good inference.
• MUST generate report.md after every training + diagnostics run. Store all artifacts in <slug>/ (see references/reporting.md for naming and structure). Save all diagnostic figures with their standard names, save metrics.csv, and produce a self-contained markdown diagnostic report. If real data is available, include the optional real-data sections in the same report. For image-valued targets, skip the standard report and use visual sample grids instead.
• NEVER use fit_online as the first training pass unless the user explicitly requests it. The first iteration MUST use fit_offline with a pre-simulated pilot budget (10k sims for fast simulators, 3k–5k for slow ones) to maximize iteration speed. Online training is a refinement step for subsequent iterations.
• MUST offer to run training in the terminal so the user can monitor progress. Training scripts should be runnable standalone; do not silently execute long training runs without giving the user access to the live output.

Installation

Install BayesFlow and a backend. Prefer JAX unless the user has a strong reason to use PyTorch or TensorFlow.

pip install "bayesflow"

BayesFlow workflow template

import bayesflow as bf
import numpy as np

# --------------------------------------------------
# 1. Define prior + observation model
# --------------------------------------------------

def my_prior():
    theta = ...
    return {"parameters": theta}

def my_observation_model(parameters):
    x = ...
    return {"observables": x}

# bf.make_simulator returns a BATCHED simulator object.
# Use simulator.sample(batch_size) or workflow.simulate(N) — NEVER loop with simulator() in Python.
simulator = bf.make_simulator([my_prior, my_observation_model])

# --------------------------------------------------
# 2. Choose architecture
# --------------------------------------------------

# See references/conditioning.md for the full conditioning logic and decision table.

# See references/model-sizes.md for different configurations — always start with Base.

# Example summary network for set-based data (exchangeable observations):
summary_net = bf.networks.SetTransformer(...) 


inference_net = bf.networks.FlowMatching(...)
# alternatives:
# bf.networks.StableConsistencyModel() # faster sampler, less performant
# bf.networks.DiffusionModel() # slower sampler, good for image generation
# bf.networks.CouplingFlow(depth=4, transform="spline") # good-old normalizing flow

# --------------------------------------------------
# 3. Build the adapter (MUST use bf.Adapter)
# --------------------------------------------------

# See references/adapter.md for the full API and a step-by-step example.
# The adapter routes simulator output to the correct network slots and handles
# parameter constraints, set assembly, dtype conversion, and concatenation.
adapter = (
    bf.Adapter()
    .as_set(["observables"])              # (N,) -> (N, 1) for SetTransformer
    .constrain("parameters", lower=0)     # if parameters have bounded support
    .convert_dtype("float64", "float32")
    .concatenate(["observables"], into="summary_variables")
    .concatenate(["parameters"], into="inference_variables")
)

# --------------------------------------------------
# 4. Create results folder and workflow
# --------------------------------------------------

import os

results_dir = "<slug>"  # e.g., "churn-model" or "churn-model-v2" for iterations
os.makedirs(results_dir, exist_ok=True)

workflow = bf.BasicWorkflow(
    simulator=simulator,
    inference_network=inference_net,
    summary_network=summary_net,
    adapter=adapter,
    checkpoint_filepath=results_dir,
)

# --------------------------------------------------
# 5. Pre-simulate pilot budget (ALWAYS offline first)
# --------------------------------------------------
# First iteration: pre-simulate a fixed budget for fast turnaround.
# - Fast simulator (< 0.05 s/draw): 20 000 datasets
# - Slow simulator (1 s+ /draw): 3 000–5 000 datasets
# Online training is a REFINEMENT step — only use it after offline
# diagnostics are healthy and you want to squeeze out more performance.

N_PILOT = 20_000  # adjust down for slow simulators
N_VAL = 300

all_sims = workflow.simulate(N_PILOT + N_VAL)

# Split into training and validation sets
train_data = {k: v[:N_PILOT] for k, v in all_sims.items()}
val_data = {k: v[N_PILOT:] for k, v in all_sims.items()}

# --------------------------------------------------
# 6. Train (offline first — fast iteration)
# --------------------------------------------------

history = workflow.fit_offline(
    data=train_data,
    epochs=100, # typically between 100 and 300
    batch_size=32,
    validation_data=val_data,
)

# --- Mandatory: save history and inspect training convergence ---
import json
from scripts.inspect_training import inspect_history

with open(os.path.join(results_dir, "history.json"), "w") as f:
    json.dump(history.history, f)

training_report = inspect_history(history.history)
print(json.dumps(training_report, indent=2))

if not training_report["overall"]["ok"]:
    print("TRAINING ISSUES — address before continuing:")
    for issue in training_report["overall"]["issues"]:
        print(f"  - {issue}")

# --------------------------------------------------
# 7. In-silico diagnostics and reporting
# --------------------------------------------------

test_data = workflow.simulate(300)

# --- Save diagnostic figures (standard names from references/reporting.md) ---
import matplotlib.pyplot as plt

figures = workflow.plot_default_diagnostics(test_data=test_data)
figure_names = {
    "losses": "loss.png",
    "recovery": "recovery.png",
    "calibration_ecdf": "calibration_ecdf.png",
    "coverage": "coverage.png",
    "z_score_contraction": "z_score_contraction.png",
}
for key, fig in figures.items():
    fig.savefig(os.path.join(results_dir, figure_names[key]), dpi=150, bbox_inches="tight")
    plt.close(fig)

# --- Save numerical diagnostics ---
metrics = workflow.compute_default_diagnostics(test_data=test_data, as_data_frame=True)
metrics.to_csv(os.path.join(results_dir, "metrics.csv"))
print(metrics)

# --- Assess and generate report ---
from scripts.check_diagnostics import check_diagnostics, suggest_next_steps

diag_report = check_diagnostics(metrics)
next_steps = suggest_next_steps(training_report, diag_report)

# Generate results_dir/report.md following the template in references/reporting.md.
# Use training_report for the Convergence assessment.
# Use diag_report["summary"] for the Recovery, Calibration, Contraction,
# and Numerical Summary assessments (plain-language ratings per parameter).
# Use next_steps for the Suggested Next Steps section.
# If real data is available, also include the optional real-data sections.

# --------------------------------------------------
# 8. Amortized inference on real data (if any)
# --------------------------------------------------

# The adapter is applied in REVERSE after sampling: the returned dict
# contains the ORIGINAL parameter names from the simulator (e.g.,
# "alpha", "beta", "sigma"), NOT "inference_variables".
# Each parameter has shape (batch, num_samples) for scalars or
# (batch, num_samples, d) for vectors.
real_data = {"observables": x_obs}
samples = workflow.sample(
    conditions=real_data,
    num_samples=1000
)
# e.g. samples["alpha"].shape == (1, 1000, 1)
#      samples["beta"].shape  == (1, 1000, 1)

# --------------------------------------------------
# 9. Posterior predictive checks (custom)
# --------------------------------------------------

# MUST reuse the existing simulator functions for PPCs.
# NEVER re-implement the generative model by hand.
# Loop over a subset of posterior draws (50 is a good default),
# indexing over the num_samples axis:
#
# n_ppc = 50
# for s in range(n_ppc):
#     # Extract single draw (index over num_samples dim)
#     theta_s = {k: samples[k][0, s] for k in ["alpha", "beta", ...]}
#     # Pass through the simulator's forward model
#     x_rep = my_observation_model(**theta_s)
#     # Compare x_rep to x_obs using domain-specific summaries

Online training (refinement step)

Online training generates fresh simulations on the fly during training. Use it only after the first offline pass shows healthy diagnostics and you want to squeeze out additional performance with a larger effective simulation budget. It is also the natural choice when the user explicitly requests it.

# Refinement: switch to online training after successful offline iteration
history = workflow.fit_online(
    epochs=200,
    batch_size=32,
    num_batches_per_epoch=100,
    validation_data=300,
)

Disk training

Use disk training when the simulation bank is too large for memory.

def custom_load(path_to_file):
    # load file + preprocessing
    return {"observables": x, "parameters": params}


workflow = bf.BasicWorkflow(
    inference_network=bf.networks.FlowMatching(...),
    summary_network=summary_net,
    inference_variables=["parameters"],
    summary_variables=["observables"],
    ...
)

history = workflow.fit_disk(root="path/to/simulation_bank", load_fn=custom_load, epochs=100, batch_size=32, validation_data=validation_data)

Augmentations

See references/augmentations.md for a guide on applying optional transformations during training.

Architecture defaults

Adapters

See references/adapter.md for the full adapter API and a step-by-step example.

When BasicWorkflow receives inference_variables=, summary_variables=, etc. as keyword arguments, it constructs a minimal implicit adapter that only renames and routes those keys. This is sufficient when the simulator output already has the right shapes and dtypes. Use an explicit bf.Adapter() chain and pass adapter= to the workflow whenever you need structural transforms (.as_set, as_time_series, .broadcast), parameter constraints (.constrain), feature engineering (.sqrt, .log), dtype coercion (.convert_dtype), or custom concatenation. The explicit adapter and the naming shorthand are mutually exclusive — do not use both.

Summary networks

See references/conditioning.md for the full conditioning model p(inference_variables | summary_net(summary_variables), inference_conditions) and a decision table.

Use a summary network (and route data through summary_variables) whenever observations are not a single fixed-length vector with meaningful element order. Most statistical models produce set-based (exchangeable) data — N i.i.d. draws, regression datasets, repeated measurements, cross-sectional samples. These MUST use a SetTransformer via summary_variables, never be flattened and placed in inference_conditions.

• Set-based / exchangeable data (most common case):

  • bf.networks.SetTransformer — required default for any model that generates N exchangeable observations
  • bf.networks.DeepSet — simpler baseline (discouraged)

• Images as conditions / observations for parameter inference:

  • bf.networks.ConvolutionalNetwork - extensible default

• Time series:

  • bf.networks.TimeSeriesNetwork — simplest default
  • bf.networks.TimeSeriesTransformer — stronger sequence model
  • bf.networks.FusionTransformer — for more complex sequential structure As a heuristic, always start by setting the summary_dim argument to 2x the number of parameters to be estimated.

• Custom summary networks — Only when the user explicitly requests one OR the data modality is non-typical (not images, time series, sets, or vectors). See references/custom-summary.md.

Inference networks

For most posterior approximation tasks, default to one of:

• bf.networks.FlowMatching() - multi-step sampling, recommended for a first iteration
• bf.networks.DiffusionModel() - multi-step sampling, recommended for high-dimensional targets
• bf.networks.StableConsistencyModel() - few-step sampling, less performant, can lose information
• bf.networks.CouplingFlow() - single-step sampling, recommended if very fast inference is important

Image-valued inference targets

If the inferential target is an image, spatial field, or other grid-valued object, use the dedicated image-generation workflow in references/image-generation.md.

• Default to bf.networks.DiffusionModel(...) with an image-capable subnet.
• Choose subnet complexity in this order: UNet < UViT < ResidualUViT.
• ConvolutionalNetwork is a summary network for image conditions, not the default choice for image-valued targets.
• Ensure all conditioning information is channel-wise concatenable with the image target; broadcast global conditions from (B, D) to (B, H, W, D) before concatenation.

Training inspection

After every fit_online, fit_offline, or fit_disk call, you must perform the following steps. Do not skip any of them.

Always pass validation_data

All three fit_* methods accept a validation_data argument. When a simulator is available, pass an integer (e.g., validation_data=300) to auto-simulate validation sets. When training offline or from disk without a simulator, pass a held-out dict of arrays. This enables val_loss tracking and overfitting detection.

Always save and inspect the returned History

fit_* returns a keras.callbacks.History object. history.history is a dict with "loss" (always present) and "val_loss" (present when validation_data is provided). BayesFlow does not save the history automatically — you must save it yourself.

After training, always:

1. Save the history to a JSON file so it survives the session:

   import json
   with open("history.json", "w") as f:
       json.dump(history.history, f)
   

2. Run scripts/inspect_training.py to get a structured convergence report:

   python scripts/inspect_training.py --history history.json
   

   Or call it in-process:

   from scripts.inspect_training import inspect_history
   report = inspect_history(history.history)
   

3. The script checks for: NaN in losses, overfitting (val_loss ratio > 1.1×), under-training (loss still decreasing), and prints a JSON report with go/no-go recommendation.

Controlling terminal output

All fit_* methods pass **kwargs through to Keras model.fit(). Use verbose=1 (default) for progress bars, verbose=2 for one line per epoch, or verbose=0 to suppress output. Prefer verbose=1 and remind the user to focus the terminal to follow how the script progresses.

Diagnostics and reporting

After every training + diagnostics run, you must generate a self-contained diagnostic report. See references/reporting.md for the full template.

Scripts

Both scripts can be run from the command line or imported as Python modules:

from scripts.inspect_training import inspect_history
from scripts.check_diagnostics import check_diagnostics, suggest_next_steps

training_report = inspect_history(history.history)
diag_report = check_diagnostics(metrics)
next_steps = suggest_next_steps(training_report, diag_report)

Diagnostic interpretation

Use workflow.compute_default_diagnostics(...) as the primary diagnostic interface. Use workflow.plot_default_diagnostics(...) as supporting visual evidence for the report.

check_diagnostics() converts numeric diagnostics into qualitative per-parameter ratings:

• calibration — rated from ECE: excellent, fair, or poor
• recovery — rated from NRMSE: excellent, good, fair, or poor
• contraction — rated from posterior contraction: high, medium, low, or poor — overconfident (high contraction + poor calibration)

The output also includes a plain-language summary per parameter (e.g., "excellent calibration; good recovery; high contraction") ready to paste into the report.

If diagnostics disagree, trust calibration first. A narrow but miscalibrated posterior is worse than a wider calibrated one.

Numeric thresholds are internal to check_diagnostics() — do not expose them in the report. Use only the qualitative ratings.

Posterior predictive checks

PPCs in BayesFlow are always custom and model-dependent. MUST reuse the existing simulator functions — NEVER re-implement the generative model by hand.

General recipe:

1. Draw posterior samples theta_s ~ q(theta | x_obs) via workflow.sample(). The returned dict has the original parameter names (e.g., alpha, beta), each with shape (batch, num_samples) or (batch, num_samples, d).
2. Loop over a subset of posterior draws (50 is a good default), indexing over the num_samples axis:

   n_ppc = 50
   for s in range(n_ppc):
       theta_s = {k: float(samples[k][0, s]) for k in ["alpha", "beta", ...]}
       x_rep = my_observation_model(**theta_s)
       # overlay / compare x_rep to x_obs
   

3. Compare x_rep to x_obs using:
   • raw overlays
   • domain-relevant summary statistics
   • discrepancy measures
   • tail behavior
   • event frequencies
   • temporal or spatial structure
4. If replicated data systematically miss the observed data, improve the simulator before trusting inference

When things go wrong

Model sizes

Always check references/model-sizes.md for rules on model sizes when choosing a summary backbone and an inference net for a particular problem.

Bayesian Workflow

Workflow overview

Every Bayesian analysis follows this sequence. Do not skip steps -- especially model criticism.

1. Formulate — Define the generative story. What underlying process, that we're precisely trying to model, created the data?
2. Specify priors — See references/priors.md
3. Implement in PyMC — Write the model. Prefer PyMC 5+ syntax. Use the latest version possible.
4. Run prior predictive checks — pm.sample_prior_predictive(). Verify priors produce plausible data ranges before fitting
5. Inference — pm.sample(nuts_sampler="nutpie"). Always use nutpie for speed (the nutpie python package provides cutting-edge sampling). Don't hardcode the number of chains — let the sampler pick the best default for the platform.
6. Diagnose convergence — Use arviz_stats.diagnose(idata) as the first check (requires arviz-stats >= 1.0.0). It covers R-hat, ESS, divergences, tree depth, and E-BFMI in one call. See references/diagnostics.md
7. Criticize the model — See references/model-criticism.md
8. Check prior sensitivity — Run psense_summary(idata) to verify conclusions are robust to prior choices. Visualize with plot_psense_dist(idata) from arviz_plots. Requires log_likelihood and log_prior in the InferenceData — compute them after sampling if needed. See references/sensitivity.md
9. Compare models (if applicable) — See references/model-comparison.md
10. Report results — Generate <slug>/report.md using the canonical template in references/reporting.md. Run scripts/check_diagnostics.py to turn raw diagnostics into qualitative ratings + an ordered next-steps list, and use that output to fill the Assessment lines and Suggested Next Steps section. When the user mentions a non-technical audience or is new to Bayesian stats, additionally adapt the prose to plain language and include a glossary — but keep the canonical report structure as the audit trail.

Installation

Prefer conda-forge / mamba-forge to install PyMC and its dependencies — pip can cause issues with compiled backends (nutpie, JAX). Example:

mamba install -c conda-forge pymc nutpie arviz arviz-stats preliz

See Stack compatibility below for the PyMC 5.x vs 6.x / ArviZ 0.23 vs 1.x notes.

Stack compatibility (PyMC 5.x and 6.x)

This skill teaches the latest PyMC 6 / ArviZ 1.x idioms and stays runnable on PyMC 5.x during the transition (regulated/corporate environments can't always upgrade freely). The scripts are verified on both stacks.

Most code is identical across versions. Where an API genuinely diverges, prefer the form that runs on both:

arviz_plots (the ArviZ 1.x plotting package, imported as azp) installs on both stacks, so leading with azp.* plots is the most portable choice. Bare az.summary(idata) works on both, but there is no single interval kwarg that does — ci_prob/ci_kind is ArviZ 1.x and hdi_prob is ArviZ 0.23 (the bare default also differs: 89% ETI vs 94% HDI). arviz_stats.diagnose, az.loo, az.plot_forest/energy/pair/khat, and idata.to_netcdf() are unchanged on both; InferenceData and DataTree are both xarray-backed, so downstream .sel() / az.* calls are identical.

PyMC model template

import pymc as pm
import arviz as az
import numpy as np

RANDOM_SEED = sum(map(ord, "churn-logistic-v1"))
rng = np.random.default_rng(RANDOM_SEED)

# always use dimensions and coordinates in PyMC models
with pm.Model(coords=coords) as model:
    # use Data containers when working on a PyMC model
    data = pm.Data("data", df["y"].to_numpy(), dims="obs")

    # --- Priors ---
    # Always document WHY each prior was chosen
    mu = pm.Normal("mu", mu=0, sigma=10)  # Weakly informative: allows wide range

    # --- Data model ---
    pm.Normal("obs", mu=mu, sigma=1, observed=data, dims="obs")

    # --- Prior predictive check ---
    prior_pred = pm.sample_prior_predictive(random_seed=rng)

    # --- Inference ---
    idata = pm.sample(nuts_sampler="nutpie", random_seed=rng)
    idata.extend(prior_pred)

    # --- Posterior predictive check ---
    idata.extend(pm.sample_posterior_predictive(idata, random_seed=rng))

    # --- Compute log-likelihood and log-prior for sensitivity checks & LOO ---
    pm.compute_log_likelihood(idata, model=model)
    pm.compute_log_prior(idata, model=model)

    # --- Save immediately after sampling ---
    # Late crashes can destroy valid results. Save to disk before any post-processing.
    idata.to_netcdf("model_output.nc")

Critical rules

• Always run prior predictive checks before sampling. If prior predictions span implausible ranges, fix priors first. If you have issues or doubts for some parameters, use the PreliZ package to elicit priors from the user.
• Always check convergence before interpreting results. R-hat > 1.01 or ESS < 100 * nbr_chains means the results are unreliable.
• Always run posterior predictive checks. A model that fits well numerically but cannot reproduce the data is useless.
• Always run calibration checks (PIT / coverage). Use ArviZ's plot_ppc_pit for this — it handles all data types (continuous, binary, count) correctly. See references/model-criticism.md.
• Document every prior choice with a brief justification in a code comment.
• Never report point estimates alone. Always include credible intervals — a 94% HDI is a fine default, but no interval width is magic (see references/reporting.md).
• Use arviz_stats.diagnose(idata) as the first diagnostic on every model (arviz-stats >= 1.0.0). It checks R-hat, ESS, divergences, tree depth saturation, and E-BFMI in one call. Follow up with az.plot_trace(idata, var_names=[...]) for visual inspection, or az.plot_rank(idata, var_names=[...]) for rank-based convergence views (both run on both stacks). Pass var_names to focus on the parameters — ArviZ 1.x errors if the auto-selected set (e.g. a vector Deterministic over an obs dim) exceeds its subplot cap. The older az.plot_trace(idata, kind="rank_vlines") is ArviZ-0.23-only.
• Don't hardcode number of chains. Let PyMC / nutpie choose the optimal default for the user's platform. Just call pm.sample() without specifying chains.
• Use reproducible, descriptive seeds. Never use magic numbers like 42. Instead, derive a seed from the analysis name: RANDOM_SEED = sum(map(ord, "my-analysis-name")). Pass it to pm.sample(random_seed=rng), pm.sample_prior_predictive(random_seed=rng), and numpy via rng = np.random.default_rng(RANDOM_SEED).
• Save InferenceData immediately after sampling with idata.to_netcdf("model_output.nc"). Late crashes or kernel restarts can destroy valid MCMC results — save before any post-processing.
• Use ArviZ for all plots and calibration. Don't write custom plotting code when ArviZ already handles it — including for binary data, count data, and calibration. ArviZ developers have thought through edge cases so you don't have to.
• Prefer xarray over numpy for InferenceData operations. InferenceData and DataTree objects are backed by xarray — use xarray's labeled indexing (.sel(), .mean(dim=...), etc.) instead of converting to numpy arrays. This preserves dimension labels, avoids shape bugs, and makes code more readable. Fall back to numpy only when xarray can't do what you need.
• Always generate <slug>/report.md after a full analysis run. Store all artifacts (inference_data.nc, trace.png, forest.png, posterior_predictive.png, pit_ecdf.png, summary.csv, diagnostics.json, calibration.json) in a slug-named results folder, and produce report.md from the canonical template in references/reporting.md. Code without an interpreted, fixed-shape report is incomplete.
• Use scripts/check_diagnostics.py to interpret diagnostics, not hand-rolled prose. Pipe the JSON outputs of diagnose_model.py (and optionally calibration_check.py and psense_summary) into check_diagnostics.py to get per-section qualitative ratings and an ordered, actionable Suggested Next Steps list. Use those outputs verbatim in the report's Assessment lines; expand only with problem-specific context.
• Always use the posterior mean (not median) for predictive probabilities. The proper Bayesian predictive distribution averages over the posterior: P(Y=k|x) = (1/S) Σ P(Y=k|x,θₛ). This is the mean, not the median. The median does not correspond to the posterior predictive distribution, can violate probability coherence (probabilities may not sum to 1), and biases calibration due to Jensen's inequality. In code: use np.mean(probs, axis=sample_axis), never np.median(...).
• Use pm.set_data() + pm.sample_posterior_predictive() for out-of-sample predictions. Don't manually extract posterior samples and recompute predictions — let PyMC propagate uncertainty properly. Define predictors as pm.Data(...) during model building, then swap in new data:

# After fitting the model:
with model:
    pm.set_data({"X": X_new, "group_idx": group_idx_new})
    oos_preds = pm.sample_posterior_predictive(idata, predictions=True, random_seed=rng)

• Check model identifiability before interpreting components. If two model components always appear together in the likelihood (e.g., a league intercept and a home advantage term when every observation is from home perspective), their individual posteriors reflect prior assumptions, not data signal — only their sum is identified. Use az.plot_pair() to check for strong posterior correlations between components. If correlation is near ±1, the components are not separately identifiable — either merge them or restructure the data.

Common model families

Utility scripts

Run these in order — each script's output feeds the next.

# 1. Run convergence + LOO + PPC checks (writes diagnostics.json)
python scripts/diagnose_model.py --idata <slug>/inference_data.nc --output <slug>/diagnostics.json

# 2. Run calibration check (writes calibration.json + pit_ecdf.png + pit_coverage.png)
python scripts/calibration_check.py --idata <slug>/inference_data.nc --output <slug>/calibration.json --save-plots --plot-dir <slug>/

# 3. Interpret the JSON outputs into qualitative ratings + suggested next steps
python scripts/check_diagnostics.py --diagnostics <slug>/diagnostics.json --calibration <slug>/calibration.json --output <slug>/check_report.json

Step 3 is what powers the report.md Assessment lines and Suggested Next Steps section — never hand-roll those interpretations from raw R-hat / ESS / pareto-k numbers when the harness can produce them consistently.

See scripts/ for all available utilities.

Common gotchas

These are battle-tested lessons that save hours of debugging:

• nutpie silently ignores idata_kwargs for log_likelihood and log_prior. Always compute them explicitly after sampling: pm.compute_log_likelihood(idata, model=model) (needed for LOO-CV) and pm.compute_log_prior(idata, model=model) (needed for prior sensitivity checks). Don't assume they're stored automatically.
• az.plot_khat() requires the LOO object, not InferenceData. Pass the output of az.loo(idata, pointwise=True) to it.
• Flat priors on scale parameters (HalfCauchy, HalfFlat) cause funnels in hierarchical models. Use Gamma(2, ...) or Exponential — these avoid the near-zero region that creates sampling problems. If there's no group-level variation to detect, you don't need the hierarchy.
• Python conditionals in models (if x > 0) don't work inside PyMC. Use pm.math.switch or pytensor.tensor.where instead.
• Forgetting to standardize predictors makes shared priors inappropriate and slows sampling. Always standardize before fitting, then back-transform for interpretation.
• Horseshoe priors create a double-funnel geometry that standard NUTS can struggle with. Always use the regularized (Finnish) horseshoe (Piironen & Vehtari, 2017), which adds a slab component that smooths the geometry. Set target_accept=0.95 or higher. If you see divergences with a horseshoe model, this is almost certainly the cause.
• np.median on posterior predictive probabilities is a silent bug. It does not produce the Bayesian predictive distribution and can yield probabilities that don't sum to 1 across categories. Always use np.mean over the posterior samples dimension.
• Discrete latents: marginalize, don't plug in. NUTS-only samplers (nutpie, Pathfinder) cannot sample discrete or ordered variables at all; PyMC's default pm.sample() falls back to a compound Metropolis step for them, which mixes poorly. Prefer to marginalize the discrete latents — pmx.marginalize(model, [...]) or analytic integration — and feed a true mixture likelihood downstream (exact, O(K) per observation, NUTS-compatible). Plugging a soft relaxation (soft-min/argmax, or E[z]) into a nonlinear function is mathematically wrong: it is not the marginal and can return out-of-bounds values. Any mixture also needs an identification constraint (e.g. ordered components) or chains will label-switch.
• Overlapping data subsets in a likelihood double-count. When a likelihood is assembled from per-subset terms, the subsets must partition the data disjointly — an observation that lands in two terms is counted twice, silently over-shrinking the posterior. Partition disjointly, or model the overlap explicitly.

When things go wrong

Causal Inference

Dependencies

This skill requires the bayesian-workflow skill for all PyMC modeling steps (priors, sampling, diagnostics, calibration, reporting).

Detect it:

ls ~/.claude/skills/bayesian-workflow/SKILL.md 2>/dev/null || ls .claude/skills/bayesian-workflow/SKILL.md 2>/dev/null

If not found, install it:

git clone https://github.com/Learning-Bayesian-Statistics/baygent-skills.git /tmp/baygent-skills
cp -r /tmp/baygent-skills/bayesian-workflow ~/.claude/skills/

For all PyMC modeling steps (priors, sampling, diagnostics, calibration, reporting), follow the bayesian-workflow skill.

Workflow overview

Every causal analysis follows this sequence. Steps 1-4 are the thinking phase (no code). Steps 5-8 are the doing phase. Think before you do.

1. Formulate the causal question — Propose precise estimand (ATE, ATT, LATE, etc.). ⚠️ ASK USER TO CONFIRM.
2. Draw the DAG — Propose causal graph with nodes, edges, and explicit non-edges. ⚠️ ASK USER TO CONFIRM. See references/dags-and-identification.md
3. Identify — Determine identification strategy (backdoor, front-door, IV, RDD, DiD). ⚠️ ASK USER TO CONFIRM untestable assumptions. See references/dags-and-identification.md
4. Choose design — Match problem to method using table below. ⚠️ ASK USER TO CONFIRM. See references/quasi-experiments.md or references/structural-models.md
5. Estimate — Build and fit the model. Delegate all PyMC mechanics to bayesian-workflow skill.
6. Refute — MANDATORY. Run design-specific robustness checks. See references/refutation.md
7. Interpret — Effect size + decision-relevant HDIs + probability of direction.
8. Report — Generate <treatment>-on-<outcome>/report.md using the canonical template in references/reporting.md. Run scripts/check_refutation.py to turn refutation outcomes into pass/marginal/fail ratings, calibrated causal language (causal / suggestive / associational / descriptive), and an ordered next-steps list. Use that output to fill the report's section 7 (causal language calibration) and Suggested Next Steps.

Design selection guide

Before reading the table, walk the intake — it routes you to a design and surfaces the assumptions you'll have to defend later:

1. Restate the estimand. ATE, ATT, LATE, or a CATE? On whom, over what period? If you can't write the one-sentence estimand, you don't have a causal question yet — you have a dataset.
2. Identify the data shape. Single time series, wide panel of units, long unit–time panel, cross-section, or pre/post group data?
3. Identify the treatment assignment. Known intervention time, staggered adoption, a threshold/cutoff, a kink, an instrument, or observed treatment with confounders?
4. State the identifying story. What makes the comparison causal — parallel trends, no anticipation, no manipulation at the cutoff, a valid instrument, overlap/positivity, donor support? This is exactly what you'll refute later.
5. Recommend one primary design + alternatives — and reconcile with the estimand. Name the best-fit design from the table plus any plausible cross-check. Then check that the estimand from step 1 is the one the design actually identifies: DiD → ATT (on the treated); RDD → a local effect at the cutoff; IV → LATE for compliers (under monotonicity); SC → the treated-unit effect. If they differ, change the design or re-scope the estimand and say so. If two designs agree later, that's convergent evidence rather than reliance on a single identification.

Critical rules

• No estimation without a confirmed DAG. A causal graph is not optional decoration — it makes assumptions explicit and determines the adjustment set. If the user resists, explain why the DAG is non-negotiable before proceeding.
• The DAG licenses adjustment, not timing — pre-treatment is necessary, not sufficient. Pre-treatment timing only rules out two post-treatment hazards: collider bias (conditioning on a common effect of treatment and outcome) and over-control (conditioning on a mediator). It does not by itself make a variable safe to adjust for. A pre-treatment variable can still be a collider on a back-door path — M-bias: a common effect of an unobserved cause of treatment and an unobserved cause of the outcome — and adjusting for it induces confounding that wasn't there. So include a covariate only if (a) it is pre-treatment and (b) the DAG shows it blocks a back-door path and is not itself a collider (or a descendant of one) on an otherwise-blocked path — i.e. it satisfies the back-door/adjustment criterion. Timing is the quick screen; the graph is the authority. (Deliberate mediation analysis is the principled exception to the post-treatment ban: it conditions on a post-treatment mediator on purpose to estimate a different estimand — direct vs. indirect effects — under stronger assumptions, including no treatment-induced mediator–outcome confounding; declare it as mediation, don't slip it in as confounder control.) If the covariates that actually satisfy the back-door criterion are too thin to identify the effect, say so and frame the result as descriptive, not causal.
• No causal claims without refutation. Every design has failure modes. Run at minimum one design-specific robustness check (placebo test, sensitivity analysis, falsification test) before reporting results. See references/refutation.md.
• State assumptions before results. Lead with what must be true for the estimate to be causal. Bury the estimate after the assumptions, not before. This is not optional politeness — it prevents misuse of results.
• Adapt HDIs to the decision context. The bayesian-workflow skill's 94% HDI is a sensible default; adapt it with explicit explanation when the decision stakes warrant it (e.g., 89% for exploratory, 97% for high-stakes policy). Report multiple intervals when the decision threshold matters.
• Downgrade causal language when warranted. If identification assumptions are unverifiable or refutation raises flags, soften claims: "consistent with a causal effect" not "causes", "estimated effect" not "true effect". Flag uncertainty loudly in the report.
• Use scripts/check_refutation.py to calibrate language, not human judgment. The script takes a structured refutation.json (see its docstring for schema) and returns a calibrated level — causal / suggestive / associational / descriptive — plus an ordered list of next steps. Use the script's output verbatim in the report; expand only with problem-specific context. Hand-rolled language calibration is exactly where overclaiming creeps in.
• Always generate <treatment>-on-<outcome>/report.md after the analysis. Store all artifacts (inference_data.nc, dag.png, effect_posterior.png, forest.png, design-specific figures, refutation.json, effect_summary.csv) in the slug-named results folder, and produce report.md from the canonical template in references/reporting.md. The report is the audit trail; code without an interpreted, fixed-shape report is incomplete.
• Ask the user when domain knowledge is needed. You cannot know whether an instrument is valid, whether parallel trends holds, or whether a confounder exists without domain expertise. Ask before assuming.
• Delegate PyMC mechanics to bayesian-workflow. This skill handles causal structure and design. The bayesian-workflow skill handles priors, sampling, diagnostics, calibration, and reporting format. Don't duplicate those rules here.

Common gotchas

These are battle-tested lessons that save hours of debugging:

• CausalPy formula syntax uses C() for categoricals. Passing a string column directly without C() will silently produce wrong dummy coding. Always wrap categorical treatment and group variables: "y ~ C(treatment) + C(group)".
• DoWhy requires explicit U nodes for unobserved confounders. Omitting them from the graph will make DoWhy treat your model as fully identified when it isn't. Add latent nodes explicitly and mark them as unobserved.
• CausalPy's PyMC models don't auto-store log-likelihood. Same issue as bayesian-workflow: nutpie silently drops it. Call pm.compute_log_likelihood(idata, model=model) after sampling if you need it for model comparison.
• Parallel trends is untestable in the post-treatment period. Pre-treatment trend tests are necessary but not sufficient — passing them doesn't prove the assumption holds after treatment. State this explicitly in every DiD report.
• Synthetic control requires the treated unit to lie within the convex hull of donors. If the treated unit is an outlier (highest GDP, largest city), no weighted combination of donors can approximate its counterfactual. Check this before running — if violated, the design is invalid.
• DiD group variable must be dummy-coded (0/1). CausalPy rejects string labels like "treatment"/"control". Use integers: 1 = treatment, 0 = control. Data also requires a unit column.
• SyntheticControl expects wide-format data. Index = time, columns = unit names, values = outcome. If your data is long format, pivot first: df.pivot(index="date", columns="unit", values="outcome").

Utility scripts

# Interpret refutation outcomes and calibrate causal language
python scripts/check_refutation.py --refutation <slug>/refutation.json --output <slug>/check_report.json

The refutation.json is written by the analyst from DoWhy/CausalPy refutation outputs. See the script docstring for the JSON schema, the test categories (critical vs. marginal), and the language calibration rules.

When things go wrong