feat: v0.9.13 — ArviZ HDI compat shim + ATT/ATU uncertainty

Small-but-load-bearing cleanup release closing two items deferred
across v0.9.10/11/12 code reviews.

Added:
- `_az_hdi_compat(samples, hdi_prob)` in statspai.bayes._base:
  routes every az.hdi() call in the Bayesian sub-package through
  one place so the inevitable arviz >= 0.18 kwarg rename
  (hdi_prob -> prob) is a one-line future change. Try hdi_prob
  first, fall back to prob on TypeError.
- ATT/ATU uncertainty on BayesianMTEResult:
  att_sd/att_hdi_lower/att_hdi_upper + atu_*. posterior_sd covers
  ATE (primary estimand), no redundant ate_sd.
- _integrated_effect returns (mean, sd, hdi_low, hdi_high); caller
  unpacks into the new fields. Six new dataclass fields all NaN-
  defaulted and appended after v0.9.12 schema (back-compat with
  pre-v0.9.13 serialised results).

Round B review: no blockers, one MEDIUM (test docstring mislabel)
fixed inline. Reviewer verified:
- Shim fallback shape correct under monkey-patched future arviz.
- Dataclass field order confirmed via live introspection.
- No __hash__ risk; .tidy()/.glance() intentionally don't surface
  the new fields (opt-in).
- Imports clean in mte.py + hte_iv.py.
- Empty-population NaN guardrail is defensive-only; unreachable
  from bayes_mte because _logit_propensity enforces 2-class
  upstream. Test renamed to reflect this.

Incident log: a regex rewrite from az.hdi(...) to _az_hdi_compat(...)
accidentally matched the helper's own body, creating a self-recursive
_az_hdi_compat. Caught by running the Bayesian suite in-session
before any test was shipped. Manual revert + re-apply.

Tests:
- tests/test_bayes_hdi_compat.py (4) — forward on current arviz,
  fall back on monkey-patched future arviz, length-2 return,
  propagate TypeError when both kwargs rejected.
- tests/test_bayes_mte_uncertainty.py (4) — SDs finite + positive,
  HDI brackets mean, no ate_sd, realistic-DGP finite.
- Bayesian family suite: 145/145 green.

Design spec: docs/superpowers/specs/2026-04-20-v0913-hdi-compat-and-att-sd.md

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

CHANGELOG.md

@@ -2,6 +2,84 @@
 
 All notable changes to StatsPAI will be documented in this file.
 
+## [0.9.13] - 2026-04-20 — ArviZ HDI compat shim + ATT/ATU uncertainty
+
+Small-but-load-bearing cleanup release. Closes two items deferred
+across the v0.9.10 / v0.9.11 / v0.9.12 code reviews.
+
+### Added (0.9.13)
+
+- **`_az_hdi_compat(samples, hdi_prob)`** in `statspai.bayes._base`
+  — calls `az.hdi(samples, hdi_prob=...)` first, falls back to
+  `az.hdi(samples, prob=...)` on `TypeError`. Routes **every**
+  `az.hdi(...)` call site in the Bayesian sub-package through one
+  place so the inevitable arviz ≥ 0.18 kwarg rename is a one-line
+  change. Previously identified as time-bomb by v0.9.12 round-C
+  review.
+
+- **ATT / ATU uncertainty** on `BayesianMTEResult`:
+  - `att_sd`, `att_hdi_lower`, `att_hdi_upper`
+  - `atu_sd`, `atu_hdi_lower`, `atu_hdi_upper`
+
+  `_integrated_effect` now returns `(mean, sd, hdi_lower,
+  hdi_upper)` instead of `(mean, sd)`. `posterior_sd` on the parent
+  result already covers ATE uncertainty — no redundant `ate_sd`.
+
+- **Appended-at-end field order** on `BayesianMTEResult` — all six
+  new fields are NaN-defaulted and positioned after the v0.9.12
+  schema (`selection`). Serialised results from earlier releases
+  deserialise cleanly.
+
+### Round-B code review found no blockers
+
+Reviewer confirmed:
+1. `_az_hdi_compat` fallback shape correct for any future arviz
+   kwarg rename.
+2. Dataclass field order verified via live introspection.
+3. No `__hash__` risk on NaN fields; broom-style `.tidy()` /
+   `.glance()` intentionally do not surface the new SD/HDI fields
+   (opt-in access).
+4. Imports clean in `mte.py` + `hte_iv.py`.
+5. Empty-population NaN guardrail is defensive-only; unreachable
+   from `bayes_mte` because `_logit_propensity` enforces 2-class
+   requirement upstream. Test renamed to reflect this honestly.
+
+One MEDIUM item (test-docstring mislabel) fixed inline.
+
+### Incident log
+
+A mass `regex` rewrite from `az.hdi(...)` to `_az_hdi_compat(...)`
+accidentally matched the helper's own body, creating a
+`_az_hdi_compat → _az_hdi_compat` self-recursion. Caught by running
+the Bayesian focused suite (would have been a stack-overflow the
+moment any Bayesian estimator shipped). Reverted + re-applied
+manually in the same session before tests ever ran outside dev.
+
+### Tests (0.9.13)
+
+- `tests/test_bayes_hdi_compat.py` (4 tests) — forwards on current
+  arviz, falls back on monkey-patched future arviz, returns length-2
+  array, propagates `TypeError` when both kwargs rejected (no silent
+  success).
+- `tests/test_bayes_mte_uncertainty.py` (4 tests) — ATT/ATU SD
+  populated + > 0, HDI brackets mean, no redundant `ate_sd`, realistic-
+  DGP both-finite.
+- Bayesian family suite: 145/145 focused MTE + sibling tests green.
+
+### Design spec
+
+- `docs/superpowers/specs/2026-04-20-v0913-hdi-compat-and-att-sd.md`
+
+### Non-goals (0.9.13)
+
+- Full bivariate-normal HV `(U_0, U_1, V) ~ N(0, Σ)` — stays queued.
+- Rust Phase 2.
+- Expose ATT/ATU HDI on `.tidy()` — today `.tidy()` describes the
+  primary estimand (ATE); adding a multi-row variant for ATT/ATU is
+  a v0.9.14+ API question.
+
+---
+
 ## [0.9.12] - 2026-04-20 — Probit-scale MTE (Heckman selection frame)
 
 Adds the third orthogonal axis to `sp.bayes_mte`: the MTE polynomial

docs/superpowers/specs/2026-04-20-v0913-hdi-compat-and-att-sd.md

@@ -0,0 +1,96 @@
+# StatsPAI v0.9.13 — arviz HDI-kwarg compat + ATT/ATU uncertainty
+
+**Author:** Bryce Wang · **Date:** 2026-04-20 · **Status:** design → implementation
+
+## 1. Motivation
+
+Two items were explicitly deferred across v0.9.10 / v0.9.11 / v0.9.12 code reviews:
+
+1. **`az.hdi(..., hdi_prob=...)` kwarg renamed to `prob` in arviz ≥ 0.18.** All our Bayesian estimators call `az.hdi(samples, hdi_prob=self.hdi_prob)`. The day a user upgrades arviz past that version, every Bayesian estimator returns malformed output. We ship a tiny compatibility shim now so the jump is a no-op.
+
+2. **`BayesianMTEResult.att` / `.atu` expose the posterior mean only.** The internal `_integrated_effect` helper already computes `(mean, sd)` but v0.9.12 discarded the SD. Without SD / HDI, users can't tell whether ATT differs from ATU meaningfully.
+
+Both are small, surgical, and high-value-per-LOC.
+
+## 2. Scope
+
+### In scope (v0.9.13)
+
+- **`_az_hdi_compat(samples, hdi_prob=...)`** in `src/statspai/bayes/_base.py` — calls `az.hdi(samples, hdi_prob=...)` first, falls back to `az.hdi(samples, prob=...)` on `TypeError`. All internal Bayesian code replaces the direct `az.hdi(...)` call with this shim.
+
+- **ATT / ATU posterior SD + HDI** on `BayesianMTEResult`:
+  - `att_sd`, `att_hdi_lower`, `att_hdi_upper`
+  - `atu_sd`, `atu_hdi_lower`, `atu_hdi_upper`
+
+- **`_integrated_effect` returns the full posterior summary** (mean, sd, hdi_low, hdi_high); caller records all four.
+
+- **Summary output** prints ATT / ATU with uncertainty when non-NaN.
+
+### Out of scope
+
+- Full bivariate-normal HV.
+- Rust Phase 2.
+- Per-unit ATT (only population-level integrated ATT).
+
+## 3. API changes
+
+### 3.1 `_az_hdi_compat`
+
+```python
+def _az_hdi_compat(samples, hdi_prob=0.95):
+    _, az = _require_pymc()
+    try:
+        return np.asarray(az.hdi(samples, hdi_prob=hdi_prob)).ravel()
+    except TypeError:
+        return np.asarray(az.hdi(samples, prob=hdi_prob)).ravel()
+```
+
+Every `az.hdi` call in `statspai.bayes.*` goes through this wrapper.
+
+### 3.2 `BayesianMTEResult` new fields
+
+Append (default NaN → preserves serialised-result backward compat):
+
+```
+att_sd, att_hdi_lower, att_hdi_upper
+atu_sd, atu_hdi_lower, atu_hdi_upper
+```
+
+`ate_sd` is NOT a new field — `posterior_sd` covers the primary estimand.
+
+### 3.3 `summary()` update
+
+When ATT/ATU SD is finite, print `ATT: 1.234 (sd 0.567, HDI [...])`.
+
+## 4. File plan
+
+| File | Change |
+|---|---|
+| `src/statspai/bayes/_base.py` | NEW `_az_hdi_compat`; `BayesianMTEResult` new fields; `summary()` extended |
+| `src/statspai/bayes/mte.py` | `_integrated_effect` returns full summary; caller wires into result |
+| `src/statspai/bayes/did.py`, `rd.py`, `iv.py`, `fuzzy_rd.py`, `hte_iv.py` | Replace `az.hdi(...)` calls with shim |
+| `tests/test_bayes_mte_uncertainty.py` | NEW |
+| `tests/test_bayes_hdi_compat.py` | NEW |
+| `pyproject.toml` | `version = "0.9.13"` |
+
+## 5. Test plan
+
+### `test_bayes_hdi_compat.py`
+
+1. `test_shim_accepts_hdi_prob_kwarg` — forwards to `az.hdi`.
+2. `test_shim_handles_typeerror` — monkey-patch `az.hdi` to raise on `hdi_prob`, confirm fallback to `prob`.
+3. `test_shim_returns_length_2_array`.
+
+### `test_bayes_mte_uncertainty.py`
+
+1. `test_att_atu_sd_fields_populated` — after fit, both SDs finite.
+2. `test_att_atu_hdi_brackets_posterior_mean` — `hdi_lower < mean < hdi_upper`.
+3. `test_summary_shows_att_atu_uncertainty` — string contains `sd ` and `HDI [`.
+4. `test_empty_population_nan_safe` — all-treated DGP → ATU SD is NaN, no crash.
+
+## 6. Success criteria
+
+1. Existing arviz still works.
+2. Mocked future-arviz (raises on `hdi_prob`) works via fallback.
+3. `r.att_sd`, `r.atu_sd` finite on a typical fit.
+4. Two rounds of code review, zero ship-blockers.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "StatsPAI"
-version = "0.9.12"
+version = "0.9.13"
 description = "The Agent-Native Causal Inference & Econometrics Toolkit for Python"
 readme = "README.md"
 license = {text = "MIT"}

src/statspai/__init__.py

@@ -22,7 +22,7 @@
 >>> sp.outreg2(result, filename="results.xlsx")
 """
 
-__version__ = "0.9.12"
+__version__ = "0.9.13"
 __author__ = "Biaoyue Wang"
 __email__ = "brycew6m@stanford.edu"
 

src/statspai/bayes/_base.py

@@ -46,6 +46,25 @@ def _require_pymc():
     return pm, az
 
 
+def _az_hdi_compat(samples, hdi_prob: float = 0.95) -> np.ndarray:
+    """Call ``arviz.hdi`` with the correct kwarg regardless of version.
+
+    arviz < 0.18 accepts ``hdi_prob=...``; arviz ≥ 0.18 renamed it to
+    ``prob=...``. Everything else downstream expects a length-2 numpy
+    array (``[lower, upper]``), so we normalise the output too.
+
+    Code-review history: the v0.9.12 round-C review flagged the
+    direct ``az.hdi(..., hdi_prob=...)`` calls as a time-bomb for
+    arviz ≥ 0.18 upgrades. This shim routes all calls through one
+    place so a future arviz rename only touches this function.
+    """
+    _, az = _require_pymc()
+    try:
+        return np.asarray(az.hdi(samples, hdi_prob=hdi_prob)).ravel()
+    except TypeError:
+        return np.asarray(az.hdi(samples, prob=hdi_prob)).ravel()
+
+
 # ---------------------------------------------------------------------------
 # Result type
 # ---------------------------------------------------------------------------
@@ -365,7 +384,7 @@ def predict_cate(self, values: Dict[str, float]) -> Dict[str, float]:
             m_vec[i] = v - self._modifier_means[i]
 
         cate_samples = tau0 + tau_hte @ m_vec
-        hdi = np.asarray(az.hdi(cate_samples, hdi_prob=self.hdi_prob)).ravel()
+        hdi = _az_hdi_compat(cate_samples, hdi_prob=self.hdi_prob)
         return {
             'mean': float(np.mean(cate_samples)),
             'median': float(np.median(cate_samples)),
@@ -396,6 +415,16 @@ class BayesianMTEResult(BayesianCausalResult):
     # remembers this so `policy_effect` can transform the abscissa
     # before integrating against the user's weight_fn.
     selection: str = 'uniform'
+    # v0.9.13: ATT / ATU uncertainty triplets. Default NaN so
+    # pre-v0.9.13 serialised results back-compat. ``posterior_sd`` on
+    # the parent already covers ATE uncertainty (ATE is the primary
+    # estimand) so there is no ``ate_sd`` field.
+    att_sd: float = float('nan')
+    att_hdi_lower: float = float('nan')
+    att_hdi_upper: float = float('nan')
+    atu_sd: float = float('nan')
+    atu_hdi_lower: float = float('nan')
+    atu_hdi_upper: float = float('nan')
 
     def policy_effect(
         self,
@@ -478,7 +507,7 @@ def policy_effect(
         numer_samples = np.trapezoid(mte_samples * weights, x=u, axis=1)
         policy_samples = numer_samples / denom
 
-        hdi = np.asarray(az.hdi(policy_samples, hdi_prob=self.hdi_prob)).ravel()
+        hdi = _az_hdi_compat(policy_samples, hdi_prob=self.hdi_prob)
         summary = {
             'label': label,
             'estimate': float(np.mean(policy_samples)),
@@ -548,9 +577,7 @@ def _summarise_posterior(
     mean = float(np.mean(posterior))
     median = float(np.median(posterior))
     sd = float(np.std(posterior, ddof=1))
-    hdi = az.hdi(posterior, hdi_prob=hdi_prob)
-    # az.hdi may return a numpy array (older arviz) or DataArray
-    hdi_arr = np.asarray(hdi).ravel()
+    hdi_arr = _az_hdi_compat(posterior, hdi_prob=hdi_prob)
     hdi_lower = float(hdi_arr[0])
     hdi_upper = float(hdi_arr[1])
     prob_positive = float(np.mean(posterior > 0))

src/statspai/bayes/hte_iv.py

@@ -22,6 +22,7 @@
 
 from ._base import (
     BayesianHTEIVResult,
+    _az_hdi_compat,
     _require_pymc,
     _sample_model,
     _summarise_posterior,
@@ -225,7 +226,7 @@ def bayes_hte_iv(
     _, az = _require_pymc()
     for k, name in enumerate(mod_cols):
         col = flat[:, k]
-        hdi = np.asarray(az.hdi(col, hdi_prob=hdi_prob)).ravel()
+        hdi = _az_hdi_compat(col, hdi_prob=hdi_prob)
         slope_rows.append({
             'term': name,
             'estimate': float(np.mean(col)),

src/statspai/bayes/mte.py

@@ -79,6 +79,7 @@
 from ._base import (
     BayesianMTEResult,
     PROBIT_CLIP,
+    _az_hdi_compat,
     _require_pymc,
     _sample_model,
     _summarise_posterior,
@@ -432,7 +433,7 @@ def _abscissa(a_expr, is_numpy: bool):
         v_grid_out = None
     for k, u in enumerate(u_grid):
         col = mte_samples[:, k]
-        hdi = np.asarray(az.hdi(col, hdi_prob=hdi_prob)).ravel()
+        hdi = _az_hdi_compat(col, hdi_prob=hdi_prob)
         row = {
             'u': float(u),
             'posterior_mean': float(np.mean(col)),
@@ -480,35 +481,41 @@ def _abscissa(a_expr, is_numpy: bool):
     U_untreated = U_pop[untreated_mask] if untreated_mask.sum() > 0 else np.array([])
 
     def _integrated_effect(U_population):
+        """Posterior summary of the MTE integrated over a subpopulation.
+
+        Returns (mean, sd, hdi_lower, hdi_high). All NaN when the
+        subpopulation is empty (e.g. all units treated ⇒ no ATU).
+        """
         if U_population.size == 0:
-            return float('nan'), float('nan')
-        # For each posterior draw, evaluate MTE at each population
-        # unit's *abscissa*. Under `selection='normal'` the polynomial
-        # is in V = Φ^{-1}(U_D), so we must transform U_population
-        # before raising to powers — otherwise ATT/ATU evaluate
-        # `b_0 + b_1·u + …` on unit-level U_D ∈ [0,1] while the
-        # posterior `b_mte` describes `b_0 + b_1·v + …` on v ∈ ℝ.
-        # That was the BLOCKER found in v0.9.12 round-B review.
+            return (float('nan'),) * 4
         if selection == 'normal':
             from scipy.stats import norm as _norm_dist
-            pop_abscissa = _norm_dist.ppf(np.clip(U_population, PROBIT_CLIP, 1 - PROBIT_CLIP))
+            pop_abscissa = _norm_dist.ppf(
+                np.clip(U_population, PROBIT_CLIP, 1 - PROBIT_CLIP)
+            )
         else:
             pop_abscissa = U_population
         u_pow_pop = np.column_stack(
             [pop_abscissa ** k for k in range(poly_u + 1)]
         )
-        samples = b_mte_post @ u_pow_pop.T  # (S, n_pop)
+        samples = b_mte_post @ u_pow_pop.T   # (S, n_pop)
         per_draw_mean = samples.mean(axis=1)  # integrated over population
-        return float(per_draw_mean.mean()), float(per_draw_mean.std(ddof=1))
+        hdi = _az_hdi_compat(per_draw_mean, hdi_prob=hdi_prob)
+        return (
+            float(per_draw_mean.mean()),
+            float(per_draw_mean.std(ddof=1)),
+            float(hdi[0]),
+            float(hdi[1]),
+        )
 
-    att_mean, _ = _integrated_effect(U_treated)
-    atu_mean, _ = _integrated_effect(U_untreated)
+    att_mean, att_sd, att_hdi_lo, att_hdi_hi = _integrated_effect(U_treated)
+    atu_mean, atu_sd, atu_hdi_lo, atu_hdi_hi = _integrated_effect(U_untreated)
 
     # Primary estimand: average MTE (ATE integral)
     ate_mean = float(ate_samples.mean())
     ate_median = float(np.median(ate_samples))
     ate_sd = float(ate_samples.std(ddof=1))
-    ate_hdi = np.asarray(az.hdi(ate_samples, hdi_prob=hdi_prob)).ravel()
+    ate_hdi = _az_hdi_compat(ate_samples, hdi_prob=hdi_prob)
     prob_pos_ate = float(np.mean(ate_samples > 0))
 
     # R-hat / ESS on the average-MTE derived quantity — not perfect
@@ -587,4 +594,10 @@ def _integrated_effect(U_population):
         att=att_mean,
         atu=atu_mean,
         selection=selection,
+        att_sd=att_sd,
+        att_hdi_lower=att_hdi_lo,
+        att_hdi_upper=att_hdi_hi,
+        atu_sd=atu_sd,
+        atu_hdi_lower=atu_hdi_lo,
+        atu_hdi_upper=atu_hdi_hi,
     )