# Pymc Inference Template

> Build a Bayesian inference project with PyMC, NUTS sampling, ArviZ diagnostics, MLflow tracking, idata persistence, and a marimo demo notebook. Use when starting any new Bayesian bundle so all bundles share the same plumbing.

- Skill: `brojonat/pymc-inference-template` (Agent Skill, multi-file: 6 files)
- Install (CLI): `npx skillmds@latest add brojonat/pymc-inference-template`
- Raw SKILL.md: https://api.skillmd.com/api/skills/brojonat/pymc-inference-template/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: AI & ML
- Author: brojonat (https://skillmd.com/u/brojonat)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/brojonat/pymc-inference-template

---


# PyMC inference template

The reference layout every Bayesian bundle in ManagerPack copies. The
goal is consistency: same project structure, same MLflow conventions,
same posterior reload path, same notebook style. Whatever the actual
model, the plumbing is identical.

The worked example is "is this coin fair?" — a Beta-Binomial conjugate
model fit with NUTS. The model is trivial, the plumbing is the point.

This is the Bayesian counterpart to `template-sklearn-pipeline`. Both
templates consume `studio/data/coin-flip.parquet`. The contrast is the
selling point: sklearn gives you a point estimate; PyMC gives you a
posterior distribution and lets you answer decision-relevant questions
like "what's the probability the coin is fair?"

## Project layout

```
<bundle>/
├── README.md             # what this bundle does + how to run it
├── SKILL.md              # this file (or specialized for the bundle)
├── src/
│   ├── train.py          # PyMC fit + ArviZ diagnostics + MLflow logging
│   ├── predict.py        # load idata from MLflow and answer questions
│   └── plots.py          # ArviZ-based plot helpers
├── notebooks/
│   └── <name>_demo.py    # marimo notebook with conjugate / interactive widgets
└── mlruns/               # MLflow tracking store (gitignored)
```

The data lives outside the bundle, in `studio/data/<problem>.parquet`,
generated by `datagen <problem>`. Bundles never carry their own data —
they consume parquet from the studio's shared data directory.

## Conventions

### Model definition

Use a `with pm.Model()` context manager. Always:

- Give every random variable an explicit name (`p`, not auto-generated)
- Document the prior choice with a comment
- Use **sufficient statistics** when they exist (e.g. `pm.Binomial("flips",
  n=N, p=p, observed=k)` rather than `N` independent Bernoulli observations).
  The likelihood is identical but the sampler is much faster.
- Note in a comment whether the model is conjugate; if so, mention the
  closed-form solution as a sanity check

```python
with pm.Model() as model:
    p = pm.Beta("p", alpha=prior_alpha, beta=prior_beta)
    pm.Binomial("flips", n=n_flips, p=p, observed=n_heads)
```

### Sampling

Always:

- Use NUTS by default (`pm.sample()` picks NUTS automatically)
- Set `random_seed=` for reproducibility
- Use 4 chains so R-hat is meaningful
- Disable progress bar in scripted runs (`progressbar=False`)
- Tune ≥ 1000, draws ≥ 2000 unless you have a reason to deviate

```python
with model:
    idata = pm.sample(
        draws=2000,
        tune=1000,
        chains=4,
        random_seed=seed,
        progressbar=False,
    )
```

### MLflow logging

Every run logs:

| Kind | What |
|---|---|
| `params` | data path, n observations, **prior parameters**, draws, tune, chains, seed, model name, sampler |
| `metrics` | posterior mean & sd, 94% HDI low/high/width, R-hat, ESS bulk/tail, decision-relevant probabilities (e.g. P(fair), P(better)), **recovery error** when ground truth is known |
| `tags` | data hash, true value tags from the sidecar |
| `artifacts` | `posterior/idata.nc` (NetCDF), `plots/posterior.png`, `plots/trace.png`, `plots/prior_vs_posterior.png`, `data/sidecar.json` |

The recovery metrics matter most:

- `posterior_mean` — point summary
- `p_recovery_error` = `|posterior_mean - true_value|`
- `true_p_in_94_hdi` — boolean, did the credible interval contain the truth?

If the truth is **always** in the HDI across many runs at the right rate
(≈94% of the time), the model is well-calibrated.

### Posterior persistence

**Save the InferenceData as a NetCDF artifact:**

```python
idata_path = tmp_dir / "idata.nc"
idata.to_netcdf(idata_path)
mlflow.log_artifact(str(idata_path), artifact_path="posterior")
```

There is no `mlflow.pymc.log_model` like there is for sklearn. The
NetCDF artifact + the script that produced it are the canonical
serialization. NetCDF is the standard format for ArviZ inference data
and round-trips cleanly.

### Loading the posterior elsewhere

```python
import arviz as az
import mlflow

mlflow.set_tracking_uri(f"file:{template_dir / 'mlruns'}")
client = mlflow.tracking.MlflowClient()
local_path = client.download_artifacts(run_id, "posterior/idata.nc")
idata = az.from_netcdf(local_path)
```

The full posterior comes back as one `InferenceData` object. No need
for the original training code.

### Diagnostics — never skip these

Always check before trusting a posterior:

- **R-hat ≤ 1.01** for every parameter (chains have mixed)
- **ESS bulk ≥ 400** per chain (effective sample size is usable)
- **No divergences** (or document why if there are some)
- **Trace plot looks like fuzzy caterpillars** (visual sanity check)

Log these as metrics so they're queryable in the MLflow UI. A bundle
that ships a model with `r_hat > 1.01` is broken.

### Plots

Use `src/plots.py` to generate ArviZ figures. Always include:

- `az.plot_posterior(idata, var_names=[...])` with the ground-truth
  marker overlaid (when known)
- `az.plot_trace(idata, var_names=[...])` for chain mixing diagnosis
- A custom prior-vs-posterior overlay using `scipy.stats` for any
  conjugate models — this makes the data's effect on the prior obvious

### Marimo notebook (always marimo, never Jupyter)

Every bundle ships a `notebooks/<name>_demo.py` marimo notebook that:

1. Loads the latest run from `mlruns/`
2. Loads the posterior via `arviz.from_netcdf`
3. Shows the posterior summary (mean, sd, HDI, R-hat, ESS)
4. **Interactive prior exploration** via `mo.ui.slider`. For conjugate
   models, use the closed-form posterior for instant updates instead of
   re-sampling. For non-conjugate models, fall back to pre-computing a
   grid of priors or accepting the slower MCMC re-fit.
5. Shows decision-relevant probabilities (e.g. P(fair))
6. Compares recovered values to ground truth from the sidecar

Validate with `marimo check notebooks/<name>_demo.py`. Test rendering
headlessly with `marimo export html notebooks/<name>_demo.py -o /tmp/x.html`
to confirm all cells execute without errors before shipping.

## Running the worked example

```bash
# 1. Generate the dataset
datagen coin-flip --n 500 --p 0.7 --seed 42

# 2. Fit and log to MLflow
cd studio/templates/pymc-inference
python src/train.py

# 3. Verify recovery
# stdout shows posterior mean ≈ 0.71, |recovery error| ≈ 0.01
# and "true p in 94% HDI: True"

# 4. Load the posterior and answer questions
python src/predict.py --run-id <id>

# 5. Open the demo notebook and explore the conjugate prior
marimo edit notebooks/coin_flip_demo.py
```

## The "is it fair?" answer

```bash
datagen coin-flip --p 0.5 --seed 99 --output /tmp/fair.parquet
python src/train.py --data /tmp/fair.parquet
# → P(p∈[0.5±0.05]) ≈ 0.6  ("60% chance the coin is fair within ±5%")

datagen coin-flip --p 0.7 --seed 42 --output /tmp/biased.parquet
python src/train.py --data /tmp/biased.parquet
# → P(p∈[0.5±0.05]) ≈ 0.0  ("definitely not fair")
```

This is the kind of decision-relevant output a sklearn point estimate
**cannot** give you. Bundles built on this template should always
expose at least one decision-relevant probability metric in the same
spot.

## Specializing this template for a new bundle

1. Copy `studio/templates/pymc-inference/` to `studio/scratch/<bundle-name>/`
   (or directly to `bundles/<bundle-name>/` once the design is settled)
2. Replace the dataset reference: change `coin-flip` to your problem name
3. Update `build_model()` to use the right priors and likelihood
4. Update the diagnostics and decision-relevant metrics
5. Update the plots in `src/plots.py` for the new domain
6. Update the marimo notebook with the right interactive widgets — for
   non-conjugate models, you may need to pre-compute a grid of priors
   instead of relying on closed-form updates
7. Update `SKILL.md` with the bundle-specific story