use super::*;
impl super::CausalAnalysis {
pub(super) fn execute_admg(
&self,
data: &TabularData,
admg: &Admg,
query: &AverageEffectQuery,
physical: &PhysicalExecutionPlan,
ctx: &ExecutionContext,
) -> Result<CausalAnalysisResult, CausalError> {
let started = Instant::now();
let identifier = physical
.logical
.record
.identifier
.as_deref()
.unwrap_or(crate::strategy_table::DEFAULT_ADMG_IDENTIFIER);
let estimator = physical
.logical
.record
.estimator
.as_deref()
.unwrap_or(crate::strategy_table::DEFAULT_ADMG_ESTIMATOR);
if !matches!(EstimatorId::parse(estimator), EstimatorId::FunctionalEffect) {
return Err(CausalError::Compile {
message: format!("ADMG ATE requires estimator functional.effect; got {estimator}"),
});
}
let identification = identify_admg(identifier, admg, query)?;
let estimand = select_estimand(&identification, EstimatorId::parse(estimator))?;
let est = FunctionalEffect {
bootstrap_replicates: self.bootstrap_replicates,
..FunctionalEffect::new()
};
let prepared = est
.prepare(
data,
&estimand,
&identification.arena,
identification.required_assumptions.clone(),
&[query.treatment, query.outcome],
)
.map_err(CausalError::from)?;
let mut ws = FunctionalDistributionWorkspace::default();
let estimate = est.estimate(&prepared, &mut ws, ctx).map_err(CausalError::from)?;
let mut diagnostics = identification.diagnostics.clone();
diagnostics.push(overlap_diagnostic(estimate.overlap));
let mut refute_ws = EstimationWorkspace::default();
let refutations = run_refuters(
data,
&estimand,
query,
&estimate,
&mut refute_ws,
None,
ctx,
self.refute,
estimator,
&self.custom_validators,
None,
)?;
let (id_artifact, id_op) = identify_provenance_step(identifier);
let (est_artifact, est_op) = estimate_provenance_step(estimator);
let provenance = provenance_pair(
(id_artifact, id_op, &[], &identification.required_assumptions),
(est_artifact, est_op, &[id_artifact], &estimate.assumptions),
);
let physical_record =
self.apply_callback_plan_marks(physical.record.clone(), &mut diagnostics);
Ok(assemble_result(AssembleArgs {
logical: &physical.logical.record,
physical: &physical_record,
identification,
estimand,
estimate,
distribution: None,
posterior: None,
mediation: None,
counterfactual: None,
anomaly: None,
change_attribution: None,
mechanism_change: None,
unit_change: None,
refutations,
diagnostics,
provenance,
treatment: query.treatment,
outcome: query.outcome,
wall_time_ns: u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX),
latency_mode: self.latency_mode.map(|m| Arc::from(m.as_str())),
stage_timings_ns: Vec::new(),
bootstrap_replicates_requested: Some(self.bootstrap_replicates),
bootstrap_replicates_ok: None,
n_draws: None,
cancelled: false,
early_stopped: false,
}))
}
pub(super) fn execute_pag(
&self,
data: &TabularData,
pag: &Pag,
query: &AverageEffectQuery,
physical: &PhysicalExecutionPlan,
ctx: &ExecutionContext,
) -> Result<CausalAnalysisResult, CausalError> {
let started = Instant::now();
let identifier = physical
.logical
.record
.identifier
.as_deref()
.unwrap_or(DEFAULT_PAG_IDENTIFIER_ID.as_str());
let estimator = physical
.logical
.record
.estimator
.as_deref()
.unwrap_or(DEFAULT_PAG_ESTIMATOR_ID.as_str());
let estimator_id = EstimatorId::parse(estimator);
let envelope = identify_pag(identifier, pag, query)?;
if matches!(envelope.status, IdentificationStatus::NotIdentified)
|| envelope.identified_weight.0 <= 0.0
{
if matches!(self.inference, InferenceMode::Bayesian(_))
|| matches!(estimator_id, EstimatorId::BayesianGcomp)
{
return self
.execute_pag_nonidentified_prior(query, physical, ctx, &envelope, started);
}
return Err(CausalError::Compile {
message: "PAG effect not identified (no identified mass in envelope)".into(),
});
}
let mut diagnostics = Vec::new();
diagnostics.push(Diagnostic::new(
"identify.pag.envelope",
DiagnosticKind::Scientific,
DiagnosticSeverity::Info,
format!(
"generalized.adjustment envelope: identified_mass={}, unidentified_mass={}, cases={}",
envelope.identified_weight.0,
envelope.unidentified_weight.0,
envelope.cases.len()
),
));
let identification = envelope_to_identification_result(&envelope, query);
if matches!(estimator_id, EstimatorId::BayesianGcomp) {
return self.execute_pag_bayesian(
data,
query,
physical,
ctx,
&envelope,
identification,
started,
);
}
let mut weighted_ate = 0.0;
let mut weighted_se2 = 0.0;
let mut total_w = 0.0;
let mut primary_estimand: Option<IdentifiedEstimand> = None;
let mut assumptions = antecedent_core::AssumptionSet::default();
for (i, case) in envelope.cases.iter().enumerate() {
if !identification_status_ok_for_case(case.result.status)
|| case.result.estimands.is_empty()
{
continue;
}
let mut estimand = select_estimand(&case.result, estimator_id.clone())?;
if estimand.method.as_ref().starts_with("generalized.adjustment") {
estimand.method = Arc::from("backdoor.adjustment");
}
let mut case_ws = StaticEstimateWorkspaces::default();
let estimate = estimate_static_effect(
estimator_id.clone(),
data,
&estimand,
query,
case.result.required_assumptions.clone(),
self.bootstrap_replicates,
self.overlap_policy,
self.population_registry.as_ref(),
ctx,
&mut case_ws,
)?;
let w = case.weight.0;
weighted_ate += w * estimate.ate;
if estimate.se_analytic.is_finite() {
weighted_se2 += w * estimate.se_analytic * estimate.se_analytic;
}
total_w += w;
if primary_estimand.is_none() {
primary_estimand = Some(estimand);
assumptions = estimate.assumptions.clone();
}
let _ = i;
}
if !matches!(total_w.partial_cmp(&0.0), Some(std::cmp::Ordering::Greater)) {
return Err(CausalError::Compile {
message: "PAG envelope had no estimable identified cases".into(),
});
}
let estimand = primary_estimand.ok_or_else(|| CausalError::Compile {
message: "PAG envelope missing estimand".into(),
})?;
let estimate = EffectEstimate::new(
weighted_ate / total_w,
(weighted_se2 / total_w).sqrt(),
assumptions.clone(),
OverlapPolicy::ExplicitOverride,
);
let mut refute_ws = EstimationWorkspace::default();
let refutations = run_refuters(
data,
&estimand,
query,
&estimate,
&mut refute_ws,
None,
ctx,
self.refute,
estimator,
&self.custom_validators,
None,
)?;
let (id_artifact, id_op) = identify_provenance_step(identifier);
let (est_artifact, est_op) = estimate_provenance_step(estimator);
let provenance = provenance_pair(
(id_artifact, id_op, &[], &identification.required_assumptions),
(est_artifact, est_op, &[id_artifact], &estimate.assumptions),
);
diagnostics.push(overlap_diagnostic(estimate.overlap));
let physical_record =
self.apply_callback_plan_marks(physical.record.clone(), &mut diagnostics);
Ok(assemble_result(AssembleArgs {
logical: &physical.logical.record,
physical: &physical_record,
identification,
estimand,
estimate,
distribution: None,
posterior: None,
mediation: None,
counterfactual: None,
anomaly: None,
change_attribution: None,
mechanism_change: None,
unit_change: None,
refutations,
diagnostics,
provenance,
treatment: query.treatment,
outcome: query.outcome,
wall_time_ns: u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX),
latency_mode: self.latency_mode.map(|m| Arc::from(m.as_str())),
stage_timings_ns: Vec::new(),
bootstrap_replicates_requested: Some(self.bootstrap_replicates),
bootstrap_replicates_ok: None,
n_draws: None,
cancelled: false,
early_stopped: false,
}))
}
pub(super) fn execute_pag_nonidentified_prior(
&self,
query: &AverageEffectQuery,
physical: &PhysicalExecutionPlan,
ctx: &ExecutionContext,
envelope: &IdentificationEnvelope<Pag>,
started: Instant,
) -> Result<CausalAnalysisResult, CausalError> {
let cfg = match &self.inference {
InferenceMode::Bayesian(c) => c.clone(),
InferenceMode::Frequentist => BayesianConfig::laplace(),
};
let scale = cfg.prior_scale.max(1e-6);
let mut prior = PriorSet::weakly_informative(1);
if let Some(g) = prior.specs.iter_mut().find_map(|s| match s {
antecedent_prob::PriorSpec::GaussianCoefficients(p) => Some(p),
_ => None,
}) {
*g = antecedent_prob::GaussianCoefficientPrior::isotropic(1, scale);
}
let posterior = nonidentified_with_prior(
&prior,
InferenceDiagnostics::analytic("pag_nonidentified_prior"),
cfg.n_draws.max(1),
ctx.rng.master_seed(),
);
let estimate = effect_from_posterior(&posterior)?;
let identification = envelope_to_identification_result(envelope, query);
let estimand = envelope.invariant.clone().unwrap_or_else(|| {
IdentifiedEstimand::backdoor(
"pag.nonidentified",
Arc::from([]),
antecedent_expr::ExprId::from_raw(0),
)
});
let mut diagnostics = identification.diagnostics.clone();
diagnostics.push(Diagnostic::new(
"estimate.pag.nonidentified_prior",
DiagnosticKind::Scientific,
DiagnosticSeverity::Warning,
format!(
"PAG not identified; returning prior-predictive draws (unidentified_mass={})",
posterior.unidentified_mass
),
));
let provenance = provenance_pair(
(
"identify.generalized_adjustment",
"identify.generalized_adjustment",
&[],
&identification.required_assumptions,
),
(
"estimate.bayesian_gcomp",
"estimate.nonidentified_with_prior",
&["identify.generalized_adjustment"],
&estimate.assumptions,
),
);
let physical_record =
self.apply_callback_plan_marks(physical.record.clone(), &mut diagnostics);
Ok(assemble_result(AssembleArgs {
logical: &physical.logical.record,
physical: &physical_record,
identification,
estimand,
estimate,
distribution: None,
posterior: Some(posterior),
mediation: None,
counterfactual: None,
anomaly: None,
change_attribution: None,
mechanism_change: None,
unit_change: None,
refutations: Vec::new(),
diagnostics,
provenance,
treatment: query.treatment,
outcome: query.outcome,
wall_time_ns: u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX),
latency_mode: self.latency_mode.map(|m| Arc::from(m.as_str())),
stage_timings_ns: Vec::new(),
bootstrap_replicates_requested: Some(self.bootstrap_replicates),
bootstrap_replicates_ok: None,
n_draws: None,
cancelled: false,
early_stopped: false,
}))
}
}