antecedent_core/query/
mediation.rs

1//! Query submodule.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use std::sync::Arc;
6
7use crate::ids::VariableId;
8use crate::intervention::Intervention;
9use crate::value::Value;
10
11use super::AverageEffectQuery;
12use super::TargetPopulation;
13use super::error::QueryError;
14
15#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
16/// Which mediation contrast to identify / estimate (linear SEM path).
17pub enum MediationContrast {
18    /// Total effect (direct + mediated).
19    Total,
20    /// Controlled / path-product direct effect (holding mediators fixed).
21    Direct,
22    /// Mediated / indirect effect (path through mediators).
23    Mediated,
24    /// Natural direct effect (linear SEM: coincides with controlled direct under linearity).
25    NaturalDirect,
26    /// Natural indirect effect (linear SEM: coincides with mediated under linearity).
27    NaturalIndirect,
28}
29
30/// Mediation query: treatment → mediators → outcome.
31#[derive(Clone, Debug, Eq, PartialEq, Hash)]
32pub struct MediationQuery {
33    /// Treatment variable.
34    pub treatment: VariableId,
35    /// Outcome variable.
36    pub outcome: VariableId,
37    /// Mediator set (non-empty).
38    pub mediators: Arc<[VariableId]>,
39    /// Contrast of interest.
40    pub contrast: MediationContrast,
41    /// Control intervention level.
42    pub control: Intervention,
43    /// Active intervention level.
44    pub active: Intervention,
45    /// Target population.
46    pub target_population: TargetPopulation,
47}
48
49impl MediationQuery {
50    /// Linear mediation with binary 0/1 treatment levels.
51    #[must_use]
52    pub fn binary(
53        treatment: VariableId,
54        outcome: VariableId,
55        mediators: impl Into<Arc<[VariableId]>>,
56        contrast: MediationContrast,
57    ) -> Self {
58        Self {
59            treatment,
60            outcome,
61            mediators: mediators.into(),
62            contrast,
63            control: Intervention::set(treatment, Value::f64(0.0)),
64            active: Intervention::set(treatment, Value::f64(1.0)),
65            target_population: TargetPopulation::AllObserved,
66        }
67    }
68
69    /// Validate ids and interventions.
70    ///
71    /// # Errors
72    ///
73    /// Empty mediators, overlaps, or inconsistent interventions.
74    pub fn validate(&self) -> Result<(), QueryError> {
75        if self.treatment == self.outcome {
76            return Err(QueryError::TreatmentEqualsOutcome { id: self.treatment });
77        }
78        if self.mediators.is_empty() {
79            return Err(QueryError::EmptyMediators);
80        }
81        if self.mediators.iter().any(|&m| m == self.treatment || m == self.outcome) {
82            return Err(QueryError::MediatorOverlapsTreatmentOrOutcome);
83        }
84        let control_var =
85            self.control.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
86        if control_var != self.treatment {
87            return Err(QueryError::InterventionVariableMismatch {
88                expected: self.treatment,
89                got: control_var,
90            });
91        }
92        let active_var =
93            self.active.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
94        if active_var != self.treatment {
95            return Err(QueryError::InterventionVariableMismatch {
96                expected: self.treatment,
97                got: active_var,
98            });
99        }
100        self.target_population.validate()?;
101        Ok(())
102    }
103}
104
105/// Conditional average effect given effect modifiers.
106#[derive(Clone, Debug, Eq, PartialEq, Hash)]
107pub struct ConditionalEffectQuery {
108    /// Inner ATE-style query; `effect_modifiers` must be non-empty.
109    pub inner: AverageEffectQuery,
110}
111
112impl ConditionalEffectQuery {
113    /// Wrap an ATE query that already carries modifiers.
114    ///
115    /// # Errors
116    ///
117    /// Empty effect modifiers.
118    pub fn try_new(inner: AverageEffectQuery) -> Result<Self, QueryError> {
119        if inner.effect_modifiers.is_empty() {
120            return Err(QueryError::EmptyEffectModifiers);
121        }
122        inner.validate()?;
123        Ok(Self { inner })
124    }
125
126    /// Validate.
127    ///
128    /// # Errors
129    ///
130    /// Empty modifiers or invalid inner query.
131    pub fn validate(&self) -> Result<(), QueryError> {
132        if self.inner.effect_modifiers.is_empty() {
133            return Err(QueryError::EmptyEffectModifiers);
134        }
135        self.inner.validate()
136    }
137}
antecedent_core/query/mediation.rs

antecedent_core/query/
mediation.rs
