antecedent_core/query/
average.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::TargetPopulation;
12use super::error::QueryError;
13
14/// Average treatment effect (ATE / ATT-style) query.
15#[derive(Clone, Debug, Eq, PartialEq, Hash)]
16#[non_exhaustive]
17pub struct AverageEffectQuery {
18    /// Treatment variable.
19    pub treatment: VariableId,
20    /// Outcome variable.
21    pub outcome: VariableId,
22    /// Optional effect modifiers .
23    pub effect_modifiers: Arc<[VariableId]>,
24    /// Control intervention level (typically treatment = 0).
25    pub control: Intervention,
26    /// Active intervention level (typically treatment = 1).
27    pub active: Intervention,
28    /// Target population.
29    pub target_population: TargetPopulation,
30}
31
32impl AverageEffectQuery {
33    /// Full constructor (required outside this crate because the type is `#[non_exhaustive]`).
34    #[must_use]
35    pub fn new(
36        treatment: VariableId,
37        outcome: VariableId,
38        effect_modifiers: impl Into<Arc<[VariableId]>>,
39        control: Intervention,
40        active: Intervention,
41        target_population: TargetPopulation,
42    ) -> Self {
43        Self {
44            treatment,
45            outcome,
46            effect_modifiers: effect_modifiers.into(),
47            control,
48            active,
49            target_population,
50        }
51    }
52
53    /// ATE for binary treatment coded as 0/1 on `treatment`.
54    ///
55    /// # Examples
56    ///
57    /// ```
58    /// use antecedent_core::{AverageEffectQuery, VariableId};
59    ///
60    /// let q = AverageEffectQuery::binary_ate(VariableId::from_raw(0), VariableId::from_raw(1));
61    /// assert_eq!(q.treatment, VariableId::from_raw(0));
62    /// ```
63    #[must_use]
64    pub fn binary_ate(treatment: VariableId, outcome: VariableId) -> Self {
65        Self::new(
66            treatment,
67            outcome,
68            Arc::from([]),
69            Intervention::set(treatment, Value::f64(0.0)),
70            Intervention::set(treatment, Value::f64(1.0)),
71            TargetPopulation::AllObserved,
72        )
73    }
74
75    /// ATE with explicit control/active float levels.
76    #[must_use]
77    pub fn with_levels(
78        treatment: VariableId,
79        outcome: VariableId,
80        control_level: f64,
81        active_level: f64,
82    ) -> Self {
83        Self {
84            treatment,
85            outcome,
86            effect_modifiers: Arc::from([]),
87            control: Intervention::set(treatment, Value::f64(control_level)),
88            active: Intervention::set(treatment, Value::f64(active_level)),
89            target_population: TargetPopulation::AllObserved,
90        }
91    }
92
93    /// Attach effect modifiers (IDs already resolved).
94    #[must_use]
95    pub fn with_effect_modifiers(mut self, modifiers: impl Into<Arc<[VariableId]>>) -> Self {
96        self.effect_modifiers = modifiers.into();
97        self
98    }
99
100    /// Set target population.
101    #[must_use]
102    pub fn with_target_population(mut self, population: TargetPopulation) -> Self {
103        self.target_population = population;
104        self
105    }
106
107    /// Validate that interventions target the treatment variable.
108    ///
109    /// # Errors
110    ///
111    /// Returns [`QueryError`] when interventions are inconsistent.
112    pub fn validate(&self) -> Result<(), QueryError> {
113        if self.treatment == self.outcome {
114            return Err(QueryError::TreatmentEqualsOutcome { id: self.treatment });
115        }
116        let control_var =
117            self.control.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
118        if control_var != self.treatment {
119            return Err(QueryError::InterventionVariableMismatch {
120                expected: self.treatment,
121                got: control_var,
122            });
123        }
124        let active_var =
125            self.active.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
126        if active_var != self.treatment {
127            return Err(QueryError::InterventionVariableMismatch {
128                expected: self.treatment,
129                got: active_var,
130            });
131        }
132        if self.effect_modifiers.iter().any(|m| *m == self.treatment || *m == self.outcome) {
133            return Err(QueryError::ModifierOverlapsTreatmentOrOutcome);
134        }
135        self.target_population.validate()?;
136        Ok(())
137    }
138}
antecedent_core/query/average.rs

antecedent_core/query/
average.rs
