antecedent_core/query/
distribution.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#[derive(Clone, Debug, PartialEq)]
15/// Interventional distribution query P(Y | do(...), Z).
16///
17/// Distinct from [`ChangeAttributionQuery`](super::ChangeAttributionQuery) (population/period change attribution).
18/// Identify via ID (empty conditioning) or IDC (nonempty conditioning);
19/// GCM sampling remains available via `sample_interventional_distribution`.
20pub struct InterventionalDistributionQuery {
21    /// Outcome variable(s) whose interventional distribution is requested.
22    pub outcomes: Arc<[VariableId]>,
23    /// Interventions defining the `do(...)` world.
24    pub interventions: Arc<[Intervention]>,
25    /// Observational conditioning set Z for `P(Y | do(X), Z)` (empty = unconditional).
26    pub conditioning: Arc<[VariableId]>,
27    /// Target population.
28    pub target_population: TargetPopulation,
29}
30
31impl InterventionalDistributionQuery {
32    /// Single-outcome interventional distribution under the given interventions.
33    #[must_use]
34    pub fn new(outcome: VariableId, interventions: impl Into<Arc<[Intervention]>>) -> Self {
35        Self {
36            outcomes: Arc::from([outcome]),
37            interventions: interventions.into(),
38            conditioning: Arc::from([]),
39            target_population: TargetPopulation::AllObserved,
40        }
41    }
42
43    /// Multiple outcomes.
44    #[must_use]
45    pub fn with_outcomes(mut self, outcomes: impl Into<Arc<[VariableId]>>) -> Self {
46        self.outcomes = outcomes.into();
47        self
48    }
49
50    /// Observational conditioning set for IDC (`P(Y | do(X), Z)`).
51    #[must_use]
52    pub fn with_conditioning(mut self, conditioning: impl Into<Arc<[VariableId]>>) -> Self {
53        self.conditioning = conditioning.into();
54        self
55    }
56
57    /// Set target population.
58    #[must_use]
59    pub fn with_target_population(mut self, population: TargetPopulation) -> Self {
60        self.target_population = population;
61        self
62    }
63
64    /// Validate outcomes, interventions, and conditioning.
65    ///
66    /// # Errors
67    ///
68    /// Empty outcomes, invalid interventions, or conditioning overlap.
69    pub fn validate(&self) -> Result<(), QueryError> {
70        if self.outcomes.is_empty() {
71            return Err(QueryError::EmptyDistributionOutcomes);
72        }
73        for iv in self.interventions.iter() {
74            iv.validate().map_err(|e| QueryError::InvalidIntervention(e.to_string()))?;
75        }
76        for &z in self.conditioning.iter() {
77            if self.outcomes.iter().any(|&y| y == z) {
78                return Err(QueryError::ConditioningOverlapsOutcomeOrIntervention);
79            }
80            if self.interventions.iter().any(|iv| iv.primary_variable() == Some(z)) {
81                return Err(QueryError::ConditioningOverlapsOutcomeOrIntervention);
82            }
83        }
84        self.target_population.validate()?;
85        Ok(())
86    }
87}
88
89/// Path-specific effect / contribution query.
90///
91/// Prefer this over overloading [`MediationQuery`](super::MediationQuery). Path *contribution*
92/// attribution is available via GCM `path_decompose`; path-restricted natural
93/// effects identify/estimate via the ID family (`path_specific.natural`) and
94/// `functional.effect` plug-in estimation.
95#[derive(Clone, Debug, PartialEq)]
96pub struct PathSpecificEffectQuery {
97    /// Treatment / source variable.
98    pub treatment: VariableId,
99    /// Outcome variable.
100    pub outcome: VariableId,
101    /// Intermediate nodes constraining the path set (`empty` = all directed paths).
102    pub path_nodes: Arc<[VariableId]>,
103    /// Control intervention level.
104    pub control: Intervention,
105    /// Active intervention level.
106    pub active: Intervention,
107    /// Target population.
108    pub target_population: TargetPopulation,
109    /// Maximum number of paths to enumerate.
110    pub max_paths: usize,
111    /// Maximum path length (edges).
112    pub max_len: usize,
113}
114
115impl PathSpecificEffectQuery {
116    /// Binary 0/1 treatment contrast with all directed paths and default limits.
117    #[must_use]
118    pub fn binary(treatment: VariableId, outcome: VariableId) -> Self {
119        Self {
120            treatment,
121            outcome,
122            path_nodes: Arc::from([]),
123            control: Intervention::set(treatment, Value::f64(0.0)),
124            active: Intervention::set(treatment, Value::f64(1.0)),
125            target_population: TargetPopulation::AllObserved,
126            max_paths: 64,
127            max_len: 16,
128        }
129    }
130
131    /// Restrict to paths that visit these intermediate nodes (in any order).
132    #[must_use]
133    pub fn with_path_nodes(mut self, nodes: impl Into<Arc<[VariableId]>>) -> Self {
134        self.path_nodes = nodes.into();
135        self
136    }
137
138    /// Cap path enumeration.
139    #[must_use]
140    pub const fn with_max_paths(mut self, max_paths: usize) -> Self {
141        self.max_paths = max_paths;
142        self
143    }
144
145    /// Cap path length.
146    #[must_use]
147    pub const fn with_max_len(mut self, max_len: usize) -> Self {
148        self.max_len = max_len;
149        self
150    }
151
152    /// Set target population.
153    #[must_use]
154    pub fn with_target_population(mut self, population: TargetPopulation) -> Self {
155        self.target_population = population;
156        self
157    }
158
159    /// Validate ids, interventions, and limits.
160    ///
161    /// # Errors
162    ///
163    /// Treatment equals outcome, zero limits, path-node overlaps, or bad interventions.
164    pub fn validate(&self) -> Result<(), QueryError> {
165        if self.treatment == self.outcome {
166            return Err(QueryError::TreatmentEqualsOutcome { id: self.treatment });
167        }
168        if self.max_paths == 0 {
169            return Err(QueryError::NonPositivePathLimit);
170        }
171        if self.max_len == 0 {
172            return Err(QueryError::NonPositivePathLimit);
173        }
174        if self.path_nodes.iter().any(|&n| n == self.treatment || n == self.outcome) {
175            return Err(QueryError::PathNodeOverlapsTreatmentOrOutcome);
176        }
177        let control_var =
178            self.control.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
179        if control_var != self.treatment {
180            return Err(QueryError::InterventionVariableMismatch {
181                expected: self.treatment,
182                got: control_var,
183            });
184        }
185        let active_var =
186            self.active.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
187        if active_var != self.treatment {
188            return Err(QueryError::InterventionVariableMismatch {
189                expected: self.treatment,
190                got: active_var,
191            });
192        }
193        self.target_population.validate()?;
194        Ok(())
195    }
196}
antecedent_core/query/distribution.rs

antecedent_core/query/
distribution.rs
