Skip to main content

antecedent_validate/
bayesian_checks.rs

1//! Prior/posterior predictive checks and prior sensitivity.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5#![allow(
6    clippy::cast_precision_loss,
7    clippy::cast_possible_truncation,
8    clippy::needless_range_loop,
9    clippy::too_many_lines,
10    clippy::many_single_char_names
11)]
12
13use std::sync::Arc;
14
15use antecedent_core::{CausalRng, ExecutionContext, KernelPolicy};
16use antecedent_estimate::{
17    BayesianGCompWorkspace, BayesianGComputationAte, CausalPosterior, PreparedBayesianProblem,
18};
19use antecedent_identify::IdentificationStatus;
20use antecedent_kernels::{PosteriorReduceOp, reduce_posterior_draws, standard_normal};
21use antecedent_prob::{
22    ExternalPriorSource, HessianFactorization, PriorSensitivitySummary, PriorSet,
23    compose_external_priors_with_alphas,
24};
25use antecedent_stats::GlmFamily;
26
27use crate::common::RefutationReport;
28use crate::error::ValidationError;
29
30/// Result of a prior or posterior predictive check.
31///
32/// Carries **two** discriminating axes so a model cannot pass merely by getting
33/// the predictive mean right: (1) location, via the mean of `mean_y` over
34/// simulations, and (2) dispersion, via the mean of the per-simulation
35/// cross-observation SD of predicted values. A model whose predictive mean is
36/// unbiased but whose predictive spread is badly wrong (e.g. off by 5×) fails
37/// on the dispersion axis even though the location axis looks fine.
38#[derive(Clone, Debug)]
39pub struct PredictiveCheckReport {
40    /// Check kind.
41    pub kind: PredictiveCheckKind,
42    /// Observed summary statistic (e.g. outcome mean).
43    pub observed: f64,
44    /// Mean of the predictive summary across simulations.
45    pub predictive_mean: f64,
46    /// SD of the predictive summary.
47    pub predictive_sd: f64,
48    /// Two-sided tail probability of `observed` under the predictive distribution.
49    pub p_value: f64,
50    /// Observed dispersion statistic: sample SD of the outcome across rows.
51    pub observed_dispersion: f64,
52    /// Mean, across simulations, of the per-simulation cross-observation SD of
53    /// predicted values (the dispersion test statistic).
54    pub predictive_dispersion_mean: f64,
55    /// Two-sided Monte Carlo p-value of `observed_dispersion` under the
56    /// simulated dispersion-statistic distribution.
57    pub dispersion_p_value: f64,
58    /// Number of predictive simulations.
59    pub n_sims: u32,
60}
61
62impl PredictiveCheckReport {
63    /// Convert to a suite [`RefutationReport`] using a two-sided α threshold on both
64    /// the location (`p_value`) and dispersion (`dispersion_p_value`) axes.
65    ///
66    /// A single mean-only statistic cannot distinguish a well-calibrated predictive
67    /// distribution from one with the right mean but the wrong spread (U/M-shaped
68    /// misspecification on variance), so both axes must clear the threshold.
69    #[must_use]
70    pub fn to_refutation_report(&self, original_ate: f64, alpha: f64) -> RefutationReport {
71        let name = match self.kind {
72            PredictiveCheckKind::Prior => "prior_predictive",
73            PredictiveCheckKind::Posterior => "posterior_predictive",
74        };
75        let mean_ok = self.p_value.is_finite() && self.p_value >= alpha;
76        let dispersion_ok = self.dispersion_p_value.is_finite() && self.dispersion_p_value >= alpha;
77        let passed = mean_ok && dispersion_ok;
78        let comparison = self.p_value.min(self.dispersion_p_value);
79        RefutationReport {
80            refuter: Arc::from(name),
81            original_ate,
82            refuted_ate: self.predictive_mean,
83            comparison,
84            informative: true,
85            passed,
86            failure_condition: if passed {
87                None
88            } else if !mean_ok && !dispersion_ok {
89                Some(Arc::from(format!(
90                    "predictive check failed on mean (p={} < alpha={alpha}) and dispersion \
91                     (p={} < alpha={alpha})",
92                    self.p_value, self.dispersion_p_value
93                )))
94            } else if !mean_ok {
95                Some(Arc::from(format!(
96                    "predictive check failed (p={} < alpha={alpha})",
97                    self.p_value
98                )))
99            } else {
100                Some(Arc::from(format!(
101                    "predictive dispersion check failed (p={} < alpha={alpha}); predictive spread \
102                     does not match observed spread even though the mean matches",
103                    self.dispersion_p_value
104                )))
105            },
106            replicates: self.n_sims,
107        }
108    }
109}
110
111/// Prior vs posterior predictive.
112#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
113pub enum PredictiveCheckKind {
114    /// Simulate from the prior predictive.
115    Prior,
116    /// Simulate from the posterior predictive.
117    Posterior,
118}
119
120/// Prior predictive check using coefficient draws from a prior (no data update)
121/// vs observed outcome mean.
122#[derive(Clone, Debug)]
123pub struct PriorPredictiveCheck {
124    /// Simulations.
125    pub n_sims: u32,
126    /// RNG seed.
127    pub seed: u64,
128    /// Mean family (inverse link applied to η before summarizing).
129    pub family: GlmFamily,
130}
131
132impl Default for PriorPredictiveCheck {
133    fn default() -> Self {
134        Self::new()
135    }
136}
137
138impl PriorPredictiveCheck {
139    /// Default 200 sims, Gaussian identity.
140    #[must_use]
141    pub fn new() -> Self {
142        Self { n_sims: 200, seed: 0, family: GlmFamily::GaussianIdentity }
143    }
144
145    /// Run against a prepared Bayesian design with a weakly informative prior.
146    ///
147    /// Prefer [`Self::check_with_prior`] when an analysis / composed prior is known.
148    ///
149    /// # Errors
150    ///
151    /// Empty design.
152    pub fn check(
153        &self,
154        problem: &PreparedBayesianProblem,
155        ctx: &ExecutionContext,
156    ) -> Result<PredictiveCheckReport, ValidationError> {
157        let p = problem.design.ncols;
158        let prior = PriorSet::weakly_informative(p);
159        self.check_with_prior(problem, &prior, ctx)
160    }
161
162    /// Run prior predictive check under an explicit coefficient prior.
163    ///
164    /// # Errors
165    ///
166    /// Empty design or missing Gaussian coefficient prior.
167    pub fn check_with_prior(
168        &self,
169        problem: &PreparedBayesianProblem,
170        prior: &PriorSet,
171        ctx: &ExecutionContext,
172    ) -> Result<PredictiveCheckReport, ValidationError> {
173        let n = problem.design.nrows;
174        let p = problem.design.ncols;
175        if n == 0 || p == 0 {
176            return Err(ValidationError::estimation_msg("empty design for PPC"));
177        }
178        let coef_prior = prior.gaussian_coefficients().ok_or_else(|| {
179            ValidationError::estimation_msg("prior missing Gaussian coefficients for PPC")
180        })?;
181        if coef_prior.len() != p {
182            return Err(ValidationError::estimation_msg(
183                "prior coefficient dimension mismatch for PPC",
184            ));
185        }
186        let mut rng = CausalRng::from_seed(self.seed);
187        let mut mean_summaries = Vec::with_capacity(self.n_sims as usize);
188        let mut disp_summaries = Vec::with_capacity(self.n_sims as usize);
189        let mut beta = vec![0.0; p];
190        let mut y_pred = vec![0.0; n];
191        for _ in 0..self.n_sims {
192            // Draw β ~ prior once per simulation, then μ_i = g^{-1}(x_i'β).
193            for c in 0..p {
194                beta[c] =
195                    coef_prior.mean[c] + coef_prior.variance[c].sqrt() * standard_normal(&mut rng);
196            }
197            for r in 0..n {
198                let mut eta = 0.0;
199                for c in 0..p {
200                    eta += problem.design.matrix[c * n + r] * beta[c];
201                }
202                y_pred[r] = self.family.mean_from_eta(eta);
203            }
204            push_mean_and_dispersion(
205                &y_pred,
206                ctx.kernel_policy,
207                &mut mean_summaries,
208                &mut disp_summaries,
209            );
210        }
211        Ok(summarize_predictive_check(
212            PredictiveCheckKind::Prior,
213            &problem.design.outcome,
214            ctx.kernel_policy,
215            &mean_summaries,
216            &disp_summaries,
217            self.n_sims,
218        ))
219    }
220}
221
222/// Posterior predictive check: resample outcome means from posterior coefficient draws.
223#[derive(Clone, Debug)]
224pub struct PosteriorPredictiveCheck {
225    /// Number of posterior draws to use (capped by available).
226    pub n_sims: u32,
227    /// Mean family (inverse link applied to η before summarizing).
228    pub family: GlmFamily,
229}
230
231impl Default for PosteriorPredictiveCheck {
232    fn default() -> Self {
233        Self::new()
234    }
235}
236
237impl PosteriorPredictiveCheck {
238    /// Default Gaussian identity.
239    #[must_use]
240    pub fn new() -> Self {
241        Self { n_sims: 200, family: GlmFamily::GaussianIdentity }
242    }
243
244    /// Check using a fitted [`CausalPosterior`] that includes coefficient columns.
245    ///
246    /// # Errors
247    ///
248    /// Missing coefficients / empty draws.
249    pub fn check(
250        &self,
251        problem: &PreparedBayesianProblem,
252        posterior: &CausalPosterior,
253    ) -> Result<PredictiveCheckReport, ValidationError> {
254        let n = problem.design.nrows;
255        let p = problem.design.ncols;
256        let n_draws = posterior.draws.n_draws.min(self.n_sims as usize);
257        if n_draws == 0 {
258            return Err(ValidationError::estimation_msg("no posterior draws for PPC"));
259        }
260        let policy = KernelPolicy::default_policy();
261        let mut mean_summaries = Vec::with_capacity(n_draws);
262        let mut disp_summaries = Vec::with_capacity(n_draws);
263        let mut y_pred = vec![0.0; n];
264        let mut beta = vec![0.0; p];
265        for d in 0..n_draws {
266            // Coefficients depend only on (d, c); hoist the fallible draw
267            // accessor out of the row loop (n×p calls per draw → p).
268            for (c, slot) in beta.iter_mut().enumerate() {
269                *slot = posterior.draws.get(d, c).map_err(ValidationError::from)?;
270            }
271            for r in 0..n {
272                let mut eta = 0.0;
273                for (c, &b) in beta.iter().enumerate() {
274                    let x = problem.design.matrix[c * n + r];
275                    eta += x * b;
276                }
277                y_pred[r] = self.family.mean_from_eta(eta);
278            }
279            push_mean_and_dispersion(&y_pred, policy, &mut mean_summaries, &mut disp_summaries);
280        }
281        Ok(summarize_predictive_check(
282            PredictiveCheckKind::Posterior,
283            &problem.design.outcome,
284            policy,
285            &mean_summaries,
286            &disp_summaries,
287            n_draws as u32,
288        ))
289    }
290}
291
292/// Default max relative range of effect means across the prior-sensitivity grid.
293pub const DEFAULT_MAX_RELATIVE_PRIOR_RANGE: f64 = 0.5;
294
295/// Prior sensitivity grid: isotropic scales **or** external α multipliers.
296#[derive(Clone, Debug)]
297pub struct PriorSensitivity {
298    /// Prior scales (σ of isotropic Gaussian coefficient prior). Empty in α mode.
299    pub scales: Arc<[f64]>,
300    /// Multipliers on post-conflict applied alphas. Empty in isotropic scale mode.
301    pub alphas: Arc<[f64]>,
302    /// Fail when `(max−min) / scale` exceeds this, where `scale` is
303    /// `max(|means…|, |original_ate|, ε)`.
304    pub max_relative_range: f64,
305}
306
307/// Inputs for external α-multiplier prior sensitivity.
308#[derive(Clone, Copy, Debug)]
309pub struct ExternalAlphaSensitivity<'a> {
310    /// Hydrated external sources (same order as composition).
311    pub sources: &'a [ExternalPriorSource],
312    /// Post-conflict applied alphas (length must match `sources`).
313    pub alphas_applied: &'a [f64],
314}
315
316impl Default for PriorSensitivity {
317    fn default() -> Self {
318        Self::standard_grid()
319    }
320}
321
322impl PriorSensitivity {
323    /// Standard isotropic grid `{0.5, 1, 2, 5, 10, 20}` with [`DEFAULT_MAX_RELATIVE_PRIOR_RANGE`].
324    #[must_use]
325    pub fn standard_grid() -> Self {
326        Self {
327            scales: Arc::from(vec![0.5, 1.0, 2.0, 5.0, 10.0, 20.0]),
328            alphas: Arc::from([]),
329            max_relative_range: DEFAULT_MAX_RELATIVE_PRIOR_RANGE,
330        }
331    }
332
333    /// Standard external-α multiplier grid `{0, 0.25, 0.5, 0.75, 1}`.
334    ///
335    /// Multiplier `0` is baseline-only; `1` uses full post-conflict applied alphas.
336    #[must_use]
337    pub fn standard_alpha_grid() -> Self {
338        Self {
339            scales: Arc::from([]),
340            alphas: Arc::from(vec![0.0, 0.25, 0.5, 0.75, 1.0]),
341            max_relative_range: DEFAULT_MAX_RELATIVE_PRIOR_RANGE,
342        }
343    }
344
345    fn grid_len(&self) -> usize {
346        if self.alphas.is_empty() { self.scales.len() } else { self.alphas.len() }
347    }
348
349    /// Refit Bayesian g-comp at each prior scale; return sensitivity summary.
350    ///
351    /// # Errors
352    ///
353    /// Fit failures or empty scale grid.
354    pub fn evaluate(
355        &self,
356        estimator: &BayesianGComputationAte,
357        problem: &PreparedBayesianProblem,
358        identification: IdentificationStatus,
359        workspace: &mut BayesianGCompWorkspace,
360        ctx: &ExecutionContext,
361    ) -> Result<(PriorSensitivitySummary, Vec<CausalPosterior>), ValidationError> {
362        if self.scales.is_empty() {
363            return Err(ValidationError::estimation_msg(
364                "prior sensitivity scale grid is empty (use evaluate_external_alpha for α mode)",
365            ));
366        }
367        let mut means = Vec::with_capacity(self.scales.len());
368        let mut sds = Vec::with_capacity(self.scales.len());
369        let mut posts = Vec::with_capacity(self.scales.len());
370        for &scale in self.scales.iter() {
371            let est = BayesianGComputationAte {
372                prior_scale: scale,
373                n_draws: estimator.n_draws.min(200),
374                seed: estimator.seed,
375                backend: estimator.backend,
376                likelihood: estimator.likelihood,
377                overlap: estimator.overlap,
378                prior: None,
379            };
380            let post = est.fit(problem, identification, workspace, ctx).map_err(|e| {
381                ValidationError::estimation_msg(format!("prior sensitivity fit failed: {e}"))
382            })?;
383            let eq = post.effect_column().ok_or_else(|| {
384                ValidationError::estimation_msg("missing effect column in sensitivity fit")
385            })?;
386            means.push(post.summaries.mean[eq]);
387            sds.push(post.summaries.sd[eq]);
388            posts.push(post);
389        }
390        Ok((
391            PriorSensitivitySummary {
392                prior_scales: Arc::clone(&self.scales),
393                alphas: Arc::from([]),
394                effect_means: Arc::from(means),
395                effect_sds: Arc::from(sds),
396            },
397            posts,
398        ))
399    }
400
401    /// Refit at each α-multiplier on post-conflict applied alphas (external prior bank).
402    ///
403    /// For multiplier `m`, composed alphas are `m * alphas_applied[k]` (clamped to `[0, 1]`).
404    ///
405    /// # Errors
406    ///
407    /// Empty α grid, length mismatch, compose failures, or fit failures.
408    pub fn evaluate_external_alpha(
409        &self,
410        estimator: &BayesianGComputationAte,
411        problem: &PreparedBayesianProblem,
412        identification: IdentificationStatus,
413        workspace: &mut BayesianGCompWorkspace,
414        ctx: &ExecutionContext,
415        external: ExternalAlphaSensitivity<'_>,
416    ) -> Result<(PriorSensitivitySummary, Vec<CausalPosterior>), ValidationError> {
417        if self.alphas.is_empty() {
418            return Err(ValidationError::estimation_msg("prior sensitivity alpha grid is empty"));
419        }
420        if external.sources.len() != external.alphas_applied.len() {
421            return Err(ValidationError::estimation_msg(
422                "evaluate_external_alpha: sources / alphas_applied length mismatch",
423            ));
424        }
425        let n_coef = problem.design.ncols;
426        let baseline = PriorSet::weakly_informative(n_coef);
427        let requested: Vec<f64> = external.sources.iter().map(|s| s.weight.alpha).collect();
428        let mut means = Vec::with_capacity(self.alphas.len());
429        let mut sds = Vec::with_capacity(self.alphas.len());
430        let mut posts = Vec::with_capacity(self.alphas.len());
431        for &mult in self.alphas.iter() {
432            if !mult.is_finite() || !(0.0..=1.0).contains(&mult) {
433                return Err(ValidationError::estimation_msg(
434                    "prior sensitivity alpha multiplier must be finite and in [0, 1]",
435                ));
436            }
437            let scaled: Vec<f64> =
438                external.alphas_applied.iter().map(|&a| (a * mult).clamp(0.0, 1.0)).collect();
439            let composed = compose_external_priors_with_alphas(
440                external.sources,
441                &requested,
442                &scaled,
443                &baseline,
444            )
445            .map_err(|e| {
446                ValidationError::estimation_msg(format!("prior sensitivity compose failed: {e}"))
447            })?;
448            let est = BayesianGComputationAte {
449                prior_scale: estimator.prior_scale,
450                n_draws: estimator.n_draws.min(200),
451                seed: estimator.seed,
452                backend: estimator.backend,
453                likelihood: estimator.likelihood,
454                overlap: estimator.overlap,
455                prior: Some(composed.prior),
456            };
457            let post = est.fit(problem, identification, workspace, ctx).map_err(|e| {
458                ValidationError::estimation_msg(format!("prior sensitivity α fit failed: {e}"))
459            })?;
460            let eq = post.effect_column().ok_or_else(|| {
461                ValidationError::estimation_msg("missing effect column in α sensitivity fit")
462            })?;
463            means.push(post.summaries.mean[eq]);
464            sds.push(post.summaries.sd[eq]);
465            posts.push(post);
466        }
467        Ok((
468            PriorSensitivitySummary {
469                prior_scales: Arc::from([]),
470                alphas: Arc::clone(&self.alphas),
471                effect_means: Arc::from(means),
472                effect_sds: Arc::from(sds),
473            },
474            posts,
475        ))
476    }
477
478    /// Convert sensitivity range into a refutation-style report.
479    ///
480    /// Passes when the relative range of effect means is finite and
481    /// `≤ max_relative_range`.
482    #[must_use]
483    pub fn to_report(
484        &self,
485        summary: &PriorSensitivitySummary,
486        original_ate: f64,
487    ) -> RefutationReport {
488        let min = summary.effect_means.iter().copied().fold(f64::INFINITY, f64::min);
489        let max = summary.effect_means.iter().copied().fold(f64::NEG_INFINITY, f64::max);
490        let range = max - min;
491        let denom = summary
492            .effect_means
493            .iter()
494            .copied()
495            .map(f64::abs)
496            .fold(original_ate.abs(), f64::max)
497            .max(1e-8);
498        let relative = range / denom;
499        let passed = relative.is_finite() && relative <= self.max_relative_range;
500        let kind =
501            if summary.alphas.is_empty() { "prior_sensitivity" } else { "prior_sensitivity_alpha" };
502        RefutationReport {
503            refuter: Arc::from(kind),
504            original_ate,
505            refuted_ate: summary.effect_means.last().copied().unwrap_or(original_ate),
506            comparison: relative,
507            informative: true,
508            passed,
509            failure_condition: if passed {
510                None
511            } else {
512                Some(Arc::from(format!(
513                    "prior sensitivity relative range {relative} exceeds max {}",
514                    self.max_relative_range
515                )))
516            },
517            replicates: u32::try_from(self.grid_len()).unwrap_or(u32::MAX),
518        }
519    }
520}
521
522fn summarize_check(
523    kind: PredictiveCheckKind,
524    observed: f64,
525    summaries: &[f64],
526    n_sims: u32,
527) -> PredictiveCheckReport {
528    let policy = KernelPolicy::default_policy();
529    let mean = reduce_posterior_draws(summaries, PosteriorReduceOp::Mean, &policy).unwrap_or(0.0);
530    let sd = reduce_posterior_draws(summaries, PosteriorReduceOp::Std, &policy).unwrap_or(0.0);
531    let n = summaries.len() as f64;
532    // (1 + count) / (1 + n) form (Davison & Hinkley): an exact-zero Monte Carlo
533    // p-value is never valid evidence with a finite sample, so both tails are
534    // bounded below by 1/(n+1) and the two-sided p-value by 2/(n+1).
535    let below = summaries.iter().filter(|&&x| x <= observed).count() as f64;
536    let above = summaries.iter().filter(|&&x| x >= observed).count() as f64;
537    let p_lower = (1.0 + below) / (1.0 + n);
538    let p_upper = (1.0 + above) / (1.0 + n);
539    let p = (2.0 * p_lower.min(p_upper)).min(1.0);
540    PredictiveCheckReport {
541        kind,
542        observed,
543        predictive_mean: mean,
544        predictive_sd: sd,
545        p_value: p,
546        // Populated by [`summarize_predictive_check`], the two-axis wrapper around this
547        // function; callers of `summarize_check` directly (unit tests exercising the
548        // Monte Carlo p-value formula) only care about the location axis above.
549        observed_dispersion: 0.0,
550        predictive_dispersion_mean: 0.0,
551        dispersion_p_value: 1.0,
552        n_sims,
553    }
554}
555
556/// Push the per-simulation location (mean) and dispersion (cross-observation SD)
557/// summary statistics for one simulated/predicted row vector `y_pred`.
558///
559/// Dispersion is the sample SD of `y_pred` *across observations within a single
560/// simulation* — how spread the predicted values are over the design — which is
561/// exactly the axis a mean-only PPC statistic is blind to (see D3 / defect C:
562/// a model with an unbiased predictive mean but a badly wrong predictive spread
563/// must be caught here, not just on the mean axis).
564fn push_mean_and_dispersion(
565    y_pred: &[f64],
566    policy: KernelPolicy,
567    mean_summaries: &mut Vec<f64>,
568    disp_summaries: &mut Vec<f64>,
569) {
570    let n = y_pred.len().max(1) as f64;
571    let mean_y = y_pred.iter().sum::<f64>() / n;
572    let sd_y = reduce_posterior_draws(y_pred, PosteriorReduceOp::Std, &policy).unwrap_or(0.0);
573    mean_summaries.push(mean_y);
574    disp_summaries.push(sd_y);
575}
576
577/// Two-axis predictive check: combines a location (mean) Monte Carlo check with a
578/// dispersion (cross-observation SD) Monte Carlo check via [`summarize_check`], and
579/// merges both into a single [`PredictiveCheckReport`].
580fn summarize_predictive_check(
581    kind: PredictiveCheckKind,
582    outcome: &[f64],
583    policy: KernelPolicy,
584    mean_summaries: &[f64],
585    disp_summaries: &[f64],
586    n_sims: u32,
587) -> PredictiveCheckReport {
588    let n = outcome.len().max(1) as f64;
589    let observed_mean = outcome.iter().sum::<f64>() / n;
590    let observed_dispersion =
591        reduce_posterior_draws(outcome, PosteriorReduceOp::Std, &policy).unwrap_or(0.0);
592    let mean_report = summarize_check(kind, observed_mean, mean_summaries, n_sims);
593    let disp_report = summarize_check(kind, observed_dispersion, disp_summaries, n_sims);
594    PredictiveCheckReport {
595        kind,
596        observed: mean_report.observed,
597        predictive_mean: mean_report.predictive_mean,
598        predictive_sd: mean_report.predictive_sd,
599        p_value: mean_report.p_value,
600        observed_dispersion,
601        predictive_dispersion_mean: disp_report.predictive_mean,
602        dispersion_p_value: disp_report.p_value,
603        n_sims,
604    }
605}
606
607/// Attach prior sensitivity onto a [`CausalPosterior`].
608#[must_use]
609pub fn with_prior_sensitivity(
610    mut posterior: CausalPosterior,
611    summary: PriorSensitivitySummary,
612) -> CausalPosterior {
613    posterior.prior_sensitivity = Some(summary);
614    posterior
615}
616
617#[cfg(test)]
618mod tests {
619    use super::*;
620    use antecedent_core::{
621        AverageEffectQuery, CausalSchemaBuilder, MeasurementSpec, RoleHint, SmallRoleSet,
622        ValueType, VariableId,
623    };
624    use antecedent_data::{
625        Float64Column, OwnedColumn, OwnedColumnarStorage, TabularData, ValidityBitmap,
626    };
627    use antecedent_estimate::{BayesianBackendKind, BayesianGComputationAte};
628    use antecedent_expr::{ExprId, IdentifiedEstimand};
629    use antecedent_identify::IdentificationStatus;
630    use antecedent_prob::{ExternalPriorWeight, GaussianCoefficientPrior, PriorSpec};
631
632    fn toy() -> (TabularData, IdentifiedEstimand, AverageEffectQuery) {
633        let n = 60usize;
634        let mut b = CausalSchemaBuilder::new();
635        b.add_variable(
636            "t",
637            ValueType::Continuous,
638            SmallRoleSet::from_hint(RoleHint::TreatmentCandidate),
639            None,
640            None,
641            MeasurementSpec::default(),
642        )
643        .unwrap();
644        b.add_variable(
645            "y",
646            ValueType::Continuous,
647            SmallRoleSet::from_hint(RoleHint::OutcomeCandidate),
648            None,
649            None,
650            MeasurementSpec::default(),
651        )
652        .unwrap();
653        b.add_variable(
654            "z",
655            ValueType::Continuous,
656            SmallRoleSet::from_hint(RoleHint::Context),
657            None,
658            None,
659            MeasurementSpec::default(),
660        )
661        .unwrap();
662        let schema = b.build().unwrap();
663        let t: Vec<f64> = (0..n).map(|i| (i % 2) as f64).collect();
664        let z: Vec<f64> = (0..n).map(|i| i as f64 * 0.05).collect();
665        let y: Vec<f64> = (0..n).map(|i| 1.0 + 2.0 * t[i] + 0.3 * z[i]).collect();
666        let cols = vec![
667            OwnedColumn::Float64(
668                Float64Column::new(
669                    VariableId::from_raw(0),
670                    Arc::from(t),
671                    ValidityBitmap::all_valid(n),
672                )
673                .unwrap(),
674            ),
675            OwnedColumn::Float64(
676                Float64Column::new(
677                    VariableId::from_raw(1),
678                    Arc::from(y),
679                    ValidityBitmap::all_valid(n),
680                )
681                .unwrap(),
682            ),
683            OwnedColumn::Float64(
684                Float64Column::new(
685                    VariableId::from_raw(2),
686                    Arc::from(z),
687                    ValidityBitmap::all_valid(n),
688                )
689                .unwrap(),
690            ),
691        ];
692        let storage = OwnedColumnarStorage::try_new(schema, cols, None, None).unwrap();
693        let estimand = IdentifiedEstimand::backdoor(
694            "backdoor.adjustment",
695            Arc::from([VariableId::from_raw(2)]),
696            ExprId::from_raw(0),
697        );
698        let query =
699            AverageEffectQuery::binary_ate(VariableId::from_raw(0), VariableId::from_raw(1));
700        (TabularData::new(storage), estimand, query)
701    }
702
703    #[test]
704    fn prior_and_posterior_ppc_run() {
705        let (data, estimand, query) = toy();
706        let bayes = BayesianGComputationAte {
707            backend: BayesianBackendKind::ConjugateGaussian,
708            n_draws: 100,
709            seed: 2,
710            prior_scale: 10.0,
711            ..BayesianGComputationAte::new()
712        };
713        let prep = bayes.prepare(&data, &estimand, &query).unwrap();
714        let ctx = ExecutionContext::for_tests(1);
715        let prior_rep = PriorPredictiveCheck { n_sims: 50, seed: 3, ..PriorPredictiveCheck::new() }
716            .check(&prep, &ctx)
717            .unwrap();
718        assert_eq!(prior_rep.kind, PredictiveCheckKind::Prior);
719        assert!(prior_rep.p_value.is_finite());
720
721        let mut ws = BayesianGCompWorkspace::default();
722        let post = bayes
723            .fit(&prep, IdentificationStatus::NonparametricallyIdentified, &mut ws, &ctx)
724            .unwrap();
725        let post_rep = PosteriorPredictiveCheck { n_sims: 50, ..PosteriorPredictiveCheck::new() }
726            .check(&prep, &post)
727            .unwrap();
728        assert_eq!(post_rep.kind, PredictiveCheckKind::Posterior);
729    }
730
731    #[test]
732    fn summarize_check_observed_outside_range_never_reports_zero() {
733        // Simulated draws are tightly clustered; an observation far outside the
734        // range on either side must not collapse the Monte Carlo p-value to
735        // exactly 0 (D1: below/n or above/n hitting 0 or 1 exactly).
736        let n = 200usize;
737        let summaries: Vec<f64> = (0..n).map(|i| i as f64 / n as f64).collect(); // [0, 1)
738        let min_p = 2.0 / (n as f64 + 1.0);
739
740        let low = summarize_check(PredictiveCheckKind::Posterior, -10.0, &summaries, n as u32);
741        assert!(low.p_value > 0.0, "p_value must be strictly positive, got {}", low.p_value);
742        assert!(low.p_value >= min_p, "p_value {} below the 2/(n+1) floor {min_p}", low.p_value);
743
744        let high = summarize_check(PredictiveCheckKind::Posterior, 10.0, &summaries, n as u32);
745        assert!(high.p_value > 0.0, "p_value must be strictly positive, got {}", high.p_value);
746        assert!(high.p_value >= min_p, "p_value {} below the 2/(n+1) floor {min_p}", high.p_value);
747    }
748
749    #[test]
750    fn summarize_check_observed_near_centre_gives_high_p_value() {
751        // Sanity check that the corrected formula still behaves as expected in
752        // the ordinary case: an observation near the middle of the simulated
753        // distribution should give a p-value near 1, not just "not exactly 0".
754        let n = 200usize;
755        let summaries: Vec<f64> = (0..n).map(|i| i as f64 / n as f64).collect(); // [0, 1)
756        let centre = summarize_check(PredictiveCheckKind::Posterior, 0.5, &summaries, n as u32);
757        assert!(centre.p_value > 0.9, "expected p_value near 1, got {}", centre.p_value);
758    }
759
760    #[test]
761    fn prior_sensitivity_grid() {
762        let (data, estimand, query) = toy();
763        let bayes = BayesianGComputationAte {
764            backend: BayesianBackendKind::ConjugateGaussian,
765            n_draws: 80,
766            seed: 4,
767            ..BayesianGComputationAte::new()
768        };
769        let prep = bayes.prepare(&data, &estimand, &query).unwrap();
770        let mut ws = BayesianGCompWorkspace::default();
771        let ctx = ExecutionContext::for_tests(1);
772        let sens = PriorSensitivity {
773            scales: Arc::from(vec![1.0, 10.0, 50.0]),
774            alphas: Arc::from([]),
775            max_relative_range: DEFAULT_MAX_RELATIVE_PRIOR_RANGE,
776        };
777        let (summary, posts) = sens
778            .evaluate(
779                &bayes,
780                &prep,
781                IdentificationStatus::NonparametricallyIdentified,
782                &mut ws,
783                &ctx,
784            )
785            .unwrap();
786        assert_eq!(summary.prior_scales.len(), 3);
787        assert!(summary.alphas.is_empty());
788        assert_eq!(posts.len(), 3);
789        let rep =
790            sens.to_report(&summary, posts[0].summaries.mean[posts[0].effect_column().unwrap()]);
791        assert!(rep.passed);
792    }
793
794    #[test]
795    fn prior_sensitivity_external_alpha_pulls_toward_source() {
796        let fixture: serde_json::Value = serde_json::from_str(include_str!(
797            "../../../conformance/validate/bayesian_checks/expected.json"
798        ))
799        .unwrap();
800        assert!(
801            fixture["contracts"]["prior_sensitivity_full_trust_moves_toward_source"]
802                .as_bool()
803                .unwrap()
804        );
805
806        let (data, estimand, query) = toy();
807        // Data ATE ≈ 2; bank a tight prior with treatment coef mean = 8.
808        let bayes = BayesianGComputationAte {
809            backend: BayesianBackendKind::ConjugateGaussian,
810            n_draws: 120,
811            seed: 7,
812            ..BayesianGComputationAte::new()
813        };
814        let prep = bayes.prepare(&data, &estimand, &query).unwrap();
815        let n = prep.design.ncols;
816        let t_col = prep.design.treatment_column().expect("treatment column");
817        let mut mean = vec![0.0; n];
818        mean[t_col] = 8.0;
819        let mut source_prior = PriorSet::new();
820        source_prior.push(PriorSpec::GaussianCoefficients(GaussianCoefficientPrior {
821            mean: Arc::from(mean),
822            variance: Arc::from(vec![0.05; n]),
823        }));
824        let sources = [ExternalPriorSource {
825            id: Arc::from("survey_a"),
826            prior: source_prior,
827            weight: ExternalPriorWeight::power(1.0).unwrap(),
828            ess: None,
829        }];
830        let alphas_applied = [1.0_f64];
831        let mut ws = BayesianGCompWorkspace::default();
832        let ctx = ExecutionContext::for_tests(1);
833        let sens = PriorSensitivity::standard_alpha_grid();
834        let (summary, _) = sens
835            .evaluate_external_alpha(
836                &bayes,
837                &prep,
838                IdentificationStatus::NonparametricallyIdentified,
839                &mut ws,
840                &ctx,
841                ExternalAlphaSensitivity { sources: &sources, alphas_applied: &alphas_applied },
842            )
843            .unwrap();
844        assert_eq!(summary.alphas.len(), 5);
845        assert!(summary.prior_scales.is_empty());
846        assert!(summary.effect_means.iter().all(|m| m.is_finite()));
847        let m0 = summary.effect_means[0];
848        let m1 = *summary.effect_means.last().unwrap();
849        // Full trust (m=1) should sit closer to the banked treatment mean than baseline (m=0).
850        assert!(
851            (m1 - 8.0).abs() < (m0 - 8.0).abs(),
852            "m=1 mean {m1} should be closer to 8 than m=0 mean {m0}"
853        );
854        let rep = sens.to_report(&summary, m1);
855        assert_eq!(rep.refuter.as_ref(), "prior_sensitivity_alpha");
856        assert!(rep.informative);
857        assert!(rep.comparison.is_finite() && rep.comparison > 0.0);
858    }
859
860    #[test]
861    fn ppc_catches_variance_misspecification_mean_ok() {
862        // Defect C regression: construct data whose outcome mean is (nearly)
863        // constant across covariates but has large residual spread (SD ~5), so a
864        // model whose predictive draws imply a small cross-observation SD gets the
865        // *mean* right (mean-only PPC would pass) while badly understating the
866        // *dispersion* of the outcome. The added dispersion axis must catch this
867        // even though the location axis does not.
868        let n = 400usize;
869        let mut b = CausalSchemaBuilder::new();
870        b.add_variable(
871            "t",
872            ValueType::Continuous,
873            SmallRoleSet::from_hint(RoleHint::TreatmentCandidate),
874            None,
875            None,
876            MeasurementSpec::default(),
877        )
878        .unwrap();
879        b.add_variable(
880            "y",
881            ValueType::Continuous,
882            SmallRoleSet::from_hint(RoleHint::OutcomeCandidate),
883            None,
884            None,
885            MeasurementSpec::default(),
886        )
887        .unwrap();
888        b.add_variable(
889            "z",
890            ValueType::Continuous,
891            SmallRoleSet::from_hint(RoleHint::Context),
892            None,
893            None,
894            MeasurementSpec::default(),
895        )
896        .unwrap();
897        let schema = b.build().unwrap();
898        // t, z are both (near-)constant so the fitted/prior mean function has almost
899        // no cross-observation spread; y has a large, fixed residual (+/- 5) on top of
900        // a tiny signal, so E[y] is well predicted but SD(y) is not.
901        let mut rng = CausalRng::from_seed(99);
902        let t: Vec<f64> = (0..n).map(|_| 0.0).collect();
903        let z: Vec<f64> = (0..n).map(|_| 0.0).collect();
904        let y: Vec<f64> = (0..n).map(|_| 5.0 * standard_normal(&mut rng)).collect();
905        let cols = vec![
906            OwnedColumn::Float64(
907                Float64Column::new(
908                    VariableId::from_raw(0),
909                    Arc::from(t),
910                    ValidityBitmap::all_valid(n),
911                )
912                .unwrap(),
913            ),
914            OwnedColumn::Float64(
915                Float64Column::new(
916                    VariableId::from_raw(1),
917                    Arc::from(y),
918                    ValidityBitmap::all_valid(n),
919                )
920                .unwrap(),
921            ),
922            OwnedColumn::Float64(
923                Float64Column::new(
924                    VariableId::from_raw(2),
925                    Arc::from(z),
926                    ValidityBitmap::all_valid(n),
927                )
928                .unwrap(),
929            ),
930        ];
931        let storage = OwnedColumnarStorage::try_new(schema, cols, None, None).unwrap();
932        let estimand = IdentifiedEstimand::backdoor(
933            "backdoor.adjustment",
934            Arc::from([VariableId::from_raw(2)]),
935            ExprId::from_raw(0),
936        );
937        let query =
938            AverageEffectQuery::binary_ate(VariableId::from_raw(0), VariableId::from_raw(1));
939        let data = TabularData::new(storage);
940
941        let bayes = BayesianGComputationAte {
942            backend: BayesianBackendKind::ConjugateGaussian,
943            n_draws: 300,
944            seed: 5,
945            // Tight prior around 0: the mean function is confidently ~0 everywhere,
946            // matching E[y]~0, but says nothing spreads out — the mean axis passes.
947            prior_scale: 0.05,
948            ..BayesianGComputationAte::new()
949        };
950        let prep = bayes.prepare(&data, &estimand, &query).unwrap();
951        let ctx = ExecutionContext::for_tests(1);
952
953        let rep = PriorPredictiveCheck { n_sims: 300, seed: 6, ..PriorPredictiveCheck::new() }
954            .check(&prep, &ctx)
955            .unwrap();
956
957        assert!(
958            rep.p_value >= 0.05,
959            "expected the mean axis to look fine (unbiased predictive mean), got p={}",
960            rep.p_value
961        );
962        assert!(
963            rep.dispersion_p_value < 0.05,
964            "expected the dispersion axis to catch the 5x variance mismatch, got \
965             dispersion_p_value={}",
966            rep.dispersion_p_value
967        );
968        let refuted = rep.to_refutation_report(0.0, 0.05);
969        assert!(
970            !refuted.passed,
971            "predictive check must fail overall when dispersion is badly wrong even \
972             though the mean matches"
973        );
974    }
975
976    #[test]
977    fn sbc_to_report_gates_on_uniformity_not_just_mean() {
978        // Defect B regression: a U-shaped rank histogram (ranks piled at the two
979        // extremes, overdispersed posterior) is symmetric about the middle bin so its
980        // mean rank fraction lands squarely in [0.35, 0.65] — a mean-only gate would
981        // pass it. The χ² uniformity statistic must catch it.
982        let n_reps = 200u32;
983        let n_draws = 100usize;
984        // Every replicate's rank is pinned at one of the two extremes (bins 0 and 9),
985        // alternating — a textbook U shape.
986        let ranks: Vec<u32> =
987            (0..n_reps).map(|i| if i % 2 == 0 { 0 } else { n_draws as u32 }).collect();
988        let n_d = n_draws as f64;
989        let fracs: Vec<f64> = ranks.iter().map(|&r| f64::from(r) / n_d).collect();
990        let mean_rank_frac = fracs.iter().sum::<f64>() / fracs.len() as f64;
991        assert!(
992            (0.35..=0.65).contains(&mean_rank_frac),
993            "fixture sanity: mean rank frac {mean_rank_frac} should look fine on its own"
994        );
995        let bins = 10usize;
996        let mut counts = vec![0.0; bins];
997        for &r in &ranks {
998            let b = ((u64::from(r) * bins as u64) / (n_draws as u64)).min(bins as u64 - 1) as usize;
999            counts[b] += 1.0;
1000        }
1001        let expected = f64::from(n_reps) / bins as f64;
1002        let mut chi2 = 0.0;
1003        for c in counts {
1004            let d = c - expected;
1005            chi2 += d * d / expected.max(1.0);
1006        }
1007        assert!(
1008            chi2 > SBC_CHI2_CRITICAL_9DF_P99,
1009            "fixture sanity: U-shaped ranks should trip the χ² statistic, got {chi2}"
1010        );
1011        let report = SbcReport { ranks: Arc::from(ranks), mean_rank_frac, uniformity_stat: chi2 };
1012        let sbc = SimulationBasedCalibration { n_reps, n_draws, seed: 0 };
1013        let rep = sbc.to_report(&report, 2.0);
1014        assert!(
1015            !rep.passed,
1016            "SBC must fail on a U-shaped (non-uniform) rank distribution even though \
1017             mean_rank_frac={mean_rank_frac:.3} is in [0.35, 0.65]"
1018        );
1019    }
1020
1021    /// Calibration-gate tests for defect A: [`SimulationBasedCalibration`] and
1022    /// [`PosteriorCalibrationOnSyntheticScm`] are fully implemented but (prior to this
1023    /// change) had zero call sites outside the `lib.rs` re-export — the Bayesian
1024    /// posterior had no coverage check at all. `#[ignore]`d for the same reason as
1025    /// `antecedent-estimate/src/calibration_coverage.rs`: many refits per test, too
1026    /// slow for every `cargo test`. Run via `cargo test -p antecedent-validate
1027    /// <name> -- --ignored --nocapture`.
1028    mod calibration_gate {
1029        use super::*;
1030
1031        fn coverage_band(n_reps: u32, level: f64) -> (f64, f64) {
1032            let se = (level * (1.0 - level) / f64::from(n_reps)).sqrt();
1033            let lo = (level - 4.0 * se).max(0.5);
1034            let hi = (level + 4.0 * se).min(1.0);
1035            (lo, hi)
1036        }
1037
1038        #[test]
1039        #[ignore = "calibration: run via scripts/gate_calibration.sh"]
1040        fn sbc_conjugate_gaussian_ranks_are_uniform() {
1041            let (data, estimand, query) = toy();
1042            let bayes = BayesianGComputationAte {
1043                backend: BayesianBackendKind::ConjugateGaussian,
1044                n_draws: 300,
1045                seed: 11,
1046                prior_scale: 5.0,
1047                ..BayesianGComputationAte::new()
1048            };
1049            let prep = bayes.prepare(&data, &estimand, &query).unwrap();
1050            let mut ws = BayesianGCompWorkspace::default();
1051            let ctx = ExecutionContext::for_tests(1);
1052            let sbc = SimulationBasedCalibration { n_reps: 200, n_draws: 300, seed: 42 };
1053            let report = sbc
1054                .check(
1055                    &bayes,
1056                    &prep,
1057                    IdentificationStatus::NonparametricallyIdentified,
1058                    &mut ws,
1059                    &ctx,
1060                )
1061                .unwrap();
1062            let rep = sbc.to_report(&report, 2.0);
1063            assert!(
1064                rep.passed,
1065                "SBC should pass for a correctly specified conjugate Gaussian model: \
1066                 mean_rank_frac={:.3} chi2={:.3}",
1067                report.mean_rank_frac, report.uniformity_stat
1068            );
1069            assert!(
1070                (0.35..=0.65).contains(&report.mean_rank_frac),
1071                "mean_rank_frac={:.3} outside [0.35, 0.65]",
1072                report.mean_rank_frac
1073            );
1074            assert!(
1075                report.uniformity_stat < SBC_CHI2_CRITICAL_9DF_P99,
1076                "chi2={:.3} exceeds critical value {SBC_CHI2_CRITICAL_9DF_P99:.3}",
1077                report.uniformity_stat
1078            );
1079        }
1080
1081        #[test]
1082        #[ignore = "calibration: run via scripts/gate_calibration.sh"]
1083        fn posterior_calibration_synthetic_scm_nominal_90_coverage() {
1084            let (data, estimand, query) = toy();
1085            let bayes = BayesianGComputationAte {
1086                backend: BayesianBackendKind::ConjugateGaussian,
1087                n_draws: 300,
1088                seed: 21,
1089                prior_scale: 5.0,
1090                ..BayesianGComputationAte::new()
1091            };
1092            let prep = bayes.prepare(&data, &estimand, &query).unwrap();
1093            let mut ws = BayesianGCompWorkspace::default();
1094            let ctx = ExecutionContext::for_tests(1);
1095            let calib = PosteriorCalibrationOnSyntheticScm {
1096                n_reps: 200,
1097                n_draws: 300,
1098                level: 0.9,
1099                seed: 77,
1100            };
1101            let report = calib
1102                .check(
1103                    &bayes,
1104                    &prep,
1105                    IdentificationStatus::NonparametricallyIdentified,
1106                    &mut ws,
1107                    &ctx,
1108                )
1109                .unwrap();
1110            let (lo, hi) = coverage_band(report.n_reps, 0.9);
1111            assert!(
1112                report.coverage >= lo && report.coverage <= hi,
1113                "nominal 90% credible-interval coverage={:.3} outside [{:.3}, {:.3}] \
1114                 ({} reps); mean_abs_error={:.3}",
1115                report.coverage,
1116                lo,
1117                hi,
1118                report.n_reps,
1119                report.mean_abs_error
1120            );
1121        }
1122    }
1123}
1124
1125/// MCMC chain diagnostics gate (ESS / R-hat / divergences).
1126///
1127/// Applicable only when the posterior was produced by an MCMC backend
1128/// (`InferenceDiagnostics::factorization == Mcmc`).
1129#[derive(Clone, Copy, Debug)]
1130pub struct McmcDiagnosticsCheck {
1131    /// Maximum acceptable split-Ř.
1132    pub max_rhat: f64,
1133    /// Minimum acceptable bulk ESS.
1134    pub min_ess: f64,
1135    /// Maximum acceptable divergence count.
1136    pub max_divergences: u32,
1137}
1138
1139impl Default for McmcDiagnosticsCheck {
1140    fn default() -> Self {
1141        Self { max_rhat: 1.05, min_ess: 10.0, max_divergences: u32::MAX / 4 }
1142    }
1143}
1144
1145impl McmcDiagnosticsCheck {
1146    /// Construct with defaults.
1147    #[must_use]
1148    pub fn new() -> Self {
1149        Self::default()
1150    }
1151
1152    /// Evaluate against a fitted posterior's diagnostics.
1153    ///
1154    /// Returns `None` when the posterior is not MCMC (caller should emit `NotApplicable`).
1155    #[must_use]
1156    pub fn check(&self, posterior: &CausalPosterior) -> Option<RefutationReport> {
1157        let d = &posterior.diagnostics;
1158        if d.factorization != HessianFactorization::Mcmc {
1159            return None;
1160        }
1161        let rhat = d.rhat_max.unwrap_or(f64::INFINITY);
1162        let ess = d.ess_bulk_min.unwrap_or(0.0);
1163        let divs = d.n_divergences.unwrap_or(u32::MAX);
1164        let passed = rhat.is_finite()
1165            && rhat <= self.max_rhat
1166            && ess >= self.min_ess
1167            && divs <= self.max_divergences
1168            && d.allows_posterior();
1169        let ate = posterior
1170            .effect_column()
1171            .and_then(|c| posterior.summaries.mean.get(c).copied())
1172            .unwrap_or(f64::NAN);
1173        Some(RefutationReport {
1174            refuter: Arc::from("mcmc_diagnostics"),
1175            original_ate: ate,
1176            refuted_ate: ate,
1177            comparison: rhat,
1178            informative: true,
1179            passed,
1180            failure_condition: if passed {
1181                None
1182            } else {
1183                Some(Arc::from(format!(
1184                    "MCMC diagnostics failed: rhat={rhat:.4} ess={ess:.1} divergences={divs}"
1185                )))
1186            },
1187            replicates: d.n_chains.unwrap_or(0),
1188        })
1189    }
1190}
1191
1192/// Simulation-based calibration ranks for a scalar posterior functional.
1193///
1194/// For each replicate: draw θ* from the prior predictive, simulate data, refit, and
1195/// record the rank of θ* among posterior draws of the primary effect.
1196#[derive(Clone, Debug)]
1197pub struct SimulationBasedCalibration {
1198    /// Number of SBC replicates.
1199    pub n_reps: u32,
1200    /// Draws per refit.
1201    pub n_draws: usize,
1202    /// RNG seed.
1203    pub seed: u64,
1204}
1205
1206impl Default for SimulationBasedCalibration {
1207    fn default() -> Self {
1208        Self { n_reps: 50, n_draws: 100, seed: 0 }
1209    }
1210}
1211
1212/// SBC report.
1213#[derive(Clone, Debug)]
1214pub struct SbcReport {
1215    /// Rank of the prior draw in each replicate (`0..=n_draws`).
1216    pub ranks: Arc<[u32]>,
1217    /// Mean rank / `n_draws` (≈ 0.5 when calibrated).
1218    pub mean_rank_frac: f64,
1219    /// Chi² uniformity diagnostic on coarse bins (lower is better).
1220    pub uniformity_stat: f64,
1221}
1222
1223impl SimulationBasedCalibration {
1224    /// Construct.
1225    #[must_use]
1226    pub fn new(n_reps: u32) -> Self {
1227        Self { n_reps: n_reps.max(1), ..Self::default() }
1228    }
1229
1230    /// Run SBC: draw θ from the prior, simulate `y` from the prior predictive under
1231    /// the fixed design matrix, refit the Bayesian g-computation estimator, and
1232    /// rank the true ATE among posterior effect draws.
1233    ///
1234    /// # Errors
1235    ///
1236    /// Fit failures.
1237    pub fn check(
1238        &self,
1239        estimator: &BayesianGComputationAte,
1240        problem: &PreparedBayesianProblem,
1241        identification: IdentificationStatus,
1242        workspace: &mut BayesianGCompWorkspace,
1243        ctx: &ExecutionContext,
1244    ) -> Result<SbcReport, ValidationError> {
1245        let mut rng = CausalRng::from_seed(self.seed);
1246        let n = problem.design.nrows;
1247        let p = problem.design.ncols;
1248        let t_col = problem
1249            .design
1250            .treatment_column()
1251            .ok_or_else(|| ValidationError::estimation_msg("SBC: missing treatment column"))?;
1252        let mut ranks = Vec::with_capacity(self.n_reps as usize);
1253        let mut est = estimator.clone();
1254        est.n_draws = self.n_draws;
1255        let scale = estimator.prior_scale.max(1e-6);
1256
1257        for rep in 0..self.n_reps {
1258            let mut beta = vec![0.0; p];
1259            for c in 0..p {
1260                beta[c] = scale * standard_normal(&mut rng);
1261            }
1262            let true_effect = (problem.active - problem.control) * beta[t_col];
1263            let mut y_rep = vec![0.0; n];
1264            for r in 0..n {
1265                let mut eta = 0.0;
1266                for c in 0..p {
1267                    eta += problem.design.matrix[c * n + r] * beta[c];
1268                }
1269                y_rep[r] = eta + standard_normal(&mut rng);
1270            }
1271            let mut sim_problem = problem.clone();
1272            let mut design = sim_problem.design.clone();
1273            design.outcome = Arc::from(y_rep);
1274            sim_problem.design = design;
1275            est.seed = self.seed ^ (u64::from(rep).wrapping_mul(0x9E37));
1276            let post = est
1277                .fit(&sim_problem, identification, workspace, ctx)
1278                .map_err(|e| ValidationError::estimation_msg(format!("SBC refit failed: {e}")))?;
1279            let col = post
1280                .effect_column()
1281                .ok_or_else(|| ValidationError::estimation_msg("SBC: no effect column"))?;
1282            let draws = post
1283                .draws
1284                .column(col)
1285                .map_err(|e| ValidationError::estimation_msg(format!("SBC draws: {e}")))?;
1286            let mut rank = 0u32;
1287            for &d in draws {
1288                if d < true_effect {
1289                    rank += 1;
1290                }
1291            }
1292            ranks.push(rank);
1293        }
1294
1295        let n_d = self.n_draws.max(1) as f64;
1296        let fracs: Vec<f64> = ranks.iter().map(|&r| f64::from(r) / n_d).collect();
1297        let mean_rank_frac =
1298            reduce_posterior_draws(&fracs, PosteriorReduceOp::Mean, &ctx.kernel_policy)
1299                .unwrap_or(0.5);
1300        let bins = 10usize;
1301        let mut counts = vec![0.0; bins];
1302        let n_draws_u = u64::try_from(self.n_draws.max(1)).unwrap_or(1);
1303        let bins_u = u64::try_from(bins).unwrap_or(1);
1304        for &r in &ranks {
1305            let b = usize::try_from(u64::from(r) * bins_u / n_draws_u).unwrap_or(0).min(bins - 1);
1306            counts[b] += 1.0;
1307        }
1308        let expected = f64::from(self.n_reps) / bins as f64;
1309        let mut chi2 = 0.0;
1310        for c in counts {
1311            let d = c - expected;
1312            chi2 += d * d / expected.max(1.0);
1313        }
1314        Ok(SbcReport { ranks: Arc::from(ranks), mean_rank_frac, uniformity_stat: chi2 })
1315    }
1316
1317    /// Convert to a refutation report.
1318    ///
1319    /// Passes only when **both** hold:
1320    /// - the mean rank fraction is in `[0.35, 0.65]` (catches gross location bias), and
1321    /// - the χ² uniformity statistic over the 10 rank bins is below
1322    ///   `SBC_CHI2_CRITICAL_9DF_P99` (catches symmetric-about-0.5 U/M-shaped rank
1323    ///   distributions — overdispersed / underdispersed posteriors — that a mean-only
1324    ///   band cannot see because they average out to ≈0.5).
1325    #[must_use]
1326    pub fn to_report(&self, report: &SbcReport, original_ate: f64) -> RefutationReport {
1327        let mean_ok = (0.35..=0.65).contains(&report.mean_rank_frac);
1328        let uniform_ok = report.uniformity_stat.is_finite()
1329            && report.uniformity_stat <= SBC_CHI2_CRITICAL_9DF_P99;
1330        let passed = mean_ok && uniform_ok;
1331        RefutationReport {
1332            refuter: Arc::from("sbc"),
1333            original_ate,
1334            refuted_ate: report.mean_rank_frac,
1335            comparison: report.uniformity_stat,
1336            informative: true,
1337            passed,
1338            failure_condition: if passed {
1339                None
1340            } else if !mean_ok && !uniform_ok {
1341                Some(Arc::from(format!(
1342                    "SBC mean rank frac {:.3} outside [0.35, 0.65] and χ²={:.3} exceeds critical \
1343                     value {SBC_CHI2_CRITICAL_9DF_P99:.3} (9 df, p=0.99)",
1344                    report.mean_rank_frac, report.uniformity_stat
1345                )))
1346            } else if !mean_ok {
1347                Some(Arc::from(format!(
1348                    "SBC mean rank frac {:.3} outside [0.35, 0.65]",
1349                    report.mean_rank_frac
1350                )))
1351            } else {
1352                Some(Arc::from(format!(
1353                    "SBC rank distribution non-uniform: χ²={:.3} exceeds critical value \
1354                     {SBC_CHI2_CRITICAL_9DF_P99:.3} (9 df, p=0.99); mean rank frac {:.3} looked \
1355                     fine but ranks are not uniformly distributed (U- or M-shaped)",
1356                    report.uniformity_stat, report.mean_rank_frac
1357                )))
1358            },
1359            replicates: self.n_reps,
1360        }
1361    }
1362}
1363
1364/// χ² critical value at the 0.99 quantile with 9 degrees of freedom (10 rank bins − 1).
1365///
1366/// Used to gate [`SimulationBasedCalibration::to_report`]'s uniformity check: SBC ranks
1367/// should be uniform under a well-calibrated posterior, and this is the standard
1368/// one-sided χ² goodness-of-fit critical value (e.g. Talts et al. 2018 §3 use the
1369/// same chi-square uniformity diagnostic).
1370const SBC_CHI2_CRITICAL_9DF_P99: f64 = 21.666;
1371
1372/// Posterior calibration on synthetic SCMs: known-ATE credible-interval coverage.
1373#[derive(Clone, Debug)]
1374pub struct PosteriorCalibrationOnSyntheticScm {
1375    /// Monte Carlo replicates.
1376    pub n_reps: u32,
1377    /// Draws per fit.
1378    pub n_draws: usize,
1379    /// Nominal coverage level (e.g. 0.9).
1380    pub level: f64,
1381    /// RNG seed.
1382    pub seed: u64,
1383}
1384
1385impl Default for PosteriorCalibrationOnSyntheticScm {
1386    fn default() -> Self {
1387        Self { n_reps: 40, n_draws: 100, level: 0.9, seed: 0 }
1388    }
1389}
1390
1391/// Report for [`PosteriorCalibrationOnSyntheticScm`].
1392#[derive(Clone, Debug)]
1393pub struct PosteriorCalibrationReport {
1394    /// Empirical coverage of equal-tailed credible intervals.
1395    pub coverage: f64,
1396    /// Mean absolute error of posterior means vs true ATE.
1397    pub mean_abs_error: f64,
1398    /// Replicates.
1399    pub n_reps: u32,
1400}
1401
1402impl PosteriorCalibrationOnSyntheticScm {
1403    /// Simulate known ATEs under the design, refit, and measure CI coverage.
1404    ///
1405    /// # Errors
1406    ///
1407    /// Fit failures.
1408    pub fn check(
1409        &self,
1410        estimator: &BayesianGComputationAte,
1411        problem: &PreparedBayesianProblem,
1412        identification: IdentificationStatus,
1413        workspace: &mut BayesianGCompWorkspace,
1414        ctx: &ExecutionContext,
1415    ) -> Result<PosteriorCalibrationReport, ValidationError> {
1416        let mut rng = CausalRng::from_seed(self.seed);
1417        let n = problem.design.nrows;
1418        let p = problem.design.ncols;
1419        let t_col = problem
1420            .design
1421            .treatment_column()
1422            .ok_or_else(|| ValidationError::estimation_msg("calibration: missing treatment"))?;
1423        let mut covered = 0u32;
1424        let mut abs_err = 0.0;
1425        let mut est = estimator.clone();
1426        est.n_draws = self.n_draws;
1427        let alpha = ((1.0 - self.level) / 2.0).clamp(0.0, 0.5);
1428
1429        for rep in 0..self.n_reps {
1430            let true_ate = standard_normal(&mut rng);
1431            let mut beta = vec![0.0; p];
1432            let diff = problem.active - problem.control;
1433            beta[t_col] = if diff.abs() > 1e-12 { true_ate / diff } else { true_ate };
1434            for c in 0..p {
1435                if c != t_col {
1436                    beta[c] = 0.5 * standard_normal(&mut rng);
1437                }
1438            }
1439            let mut y = vec![0.0; n];
1440            for r in 0..n {
1441                let mut eta = 0.0;
1442                for c in 0..p {
1443                    eta += problem.design.matrix[c * n + r] * beta[c];
1444                }
1445                y[r] = eta + standard_normal(&mut rng);
1446            }
1447            let mut sim = problem.clone();
1448            let mut design = sim.design.clone();
1449            design.outcome = Arc::from(y);
1450            sim.design = design;
1451            est.seed = self.seed ^ (u64::from(rep).wrapping_mul(0xC2B2));
1452            let post = est
1453                .fit(&sim, identification, workspace, ctx)
1454                .map_err(|e| ValidationError::estimation_msg(format!("calibration refit: {e}")))?;
1455            let col = post
1456                .effect_column()
1457                .ok_or_else(|| ValidationError::estimation_msg("calibration: no effect"))?;
1458            let mut draws = post.draws.column(col).map_err(ValidationError::from)?.to_vec();
1459            draws.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1460            let lo = quantile_sorted(&draws, alpha);
1461            let hi = quantile_sorted(&draws, 1.0 - alpha);
1462            let mean = reduce_posterior_draws(&draws, PosteriorReduceOp::Mean, &ctx.kernel_policy)
1463                .unwrap_or(0.0);
1464            abs_err += (mean - true_ate).abs();
1465            if true_ate >= lo && true_ate <= hi {
1466                covered += 1;
1467            }
1468        }
1469        Ok(PosteriorCalibrationReport {
1470            coverage: f64::from(covered) / f64::from(self.n_reps.max(1)),
1471            mean_abs_error: abs_err / f64::from(self.n_reps.max(1)),
1472            n_reps: self.n_reps,
1473        })
1474    }
1475}
1476
1477fn quantile_sorted(sorted: &[f64], q: f64) -> f64 {
1478    if sorted.is_empty() {
1479        return 0.0;
1480    }
1481    let max_idx = sorted.len() - 1;
1482    let rank = (max_idx as f64 * q.clamp(0.0, 1.0)).round();
1483    let idx = (0..=max_idx)
1484        .min_by(|&a, &b| {
1485            (a as f64 - rank)
1486                .abs()
1487                .partial_cmp(&(b as f64 - rank).abs())
1488                .unwrap_or(std::cmp::Ordering::Equal)
1489        })
1490        .unwrap_or(0);
1491    sorted[idx]
1492}