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