antecedent_core/query/
temporal.rs

1//! Query submodule.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use crate::ids::VariableId;
6use crate::intervention::{Intervention, TemporalPolicy};
7use crate::value::Value;
8
9use super::TargetPopulation;
10use super::error::QueryError;
11
12/// Temporal effect query over a discrete horizon.
13#[derive(Clone, Debug, PartialEq)]
14pub struct TemporalEffectQuery {
15    /// Treatment variable.
16    pub treatment: VariableId,
17    /// Outcome variable.
18    pub outcome: VariableId,
19    /// Temporal intervention policy.
20    pub policy: TemporalPolicy,
21    /// Control intervention level on the treatment variable.
22    pub control: Intervention,
23    /// Active intervention level on the treatment variable.
24    pub active: Intervention,
25    /// Outcome horizon in time steps after the policy origin (must be ≥ 1).
26    pub horizon_steps: u32,
27    /// Optional max history lag (steps) to retain when unfolding; `None` = planner default.
28    pub max_history_lag: Option<u32>,
29    /// Target population.
30    pub target_population: TargetPopulation,
31}
32
33impl TemporalEffectQuery {
34    /// Pulse intervention at step 0 with active float level; control is 0.0.
35    #[must_use]
36    pub fn pulse(treatment: VariableId, outcome: VariableId, active_level: f64) -> Self {
37        Self {
38            treatment,
39            outcome,
40            policy: TemporalPolicy::pulse(0),
41            control: Intervention::set(treatment, Value::f64(0.0)),
42            active: Intervention::set(treatment, Value::f64(active_level)),
43            horizon_steps: 1,
44            max_history_lag: None,
45            target_population: TargetPopulation::AllObserved,
46        }
47    }
48
49    /// Sustained intervention on `[0, until]` with active float level; control is 0.0.
50    #[must_use]
51    pub fn sustained(
52        treatment: VariableId,
53        outcome: VariableId,
54        until: i32,
55        active_level: f64,
56    ) -> Self {
57        Self {
58            treatment,
59            outcome,
60            policy: TemporalPolicy::sustained(0, until),
61            control: Intervention::set(treatment, Value::f64(0.0)),
62            active: Intervention::set(treatment, Value::f64(active_level)),
63            horizon_steps: 1,
64            max_history_lag: None,
65            target_population: TargetPopulation::AllObserved,
66        }
67    }
68
69    /// Set outcome evaluation horizon in time steps.
70    #[must_use]
71    pub const fn with_horizon_steps(mut self, horizon_steps: u32) -> Self {
72        self.horizon_steps = horizon_steps;
73        self
74    }
75
76    /// Set optional max history lag for unfolding.
77    #[must_use]
78    pub const fn with_max_history_lag(mut self, max_history_lag: Option<u32>) -> Self {
79        self.max_history_lag = max_history_lag;
80        self
81    }
82
83    /// Replace the temporal policy.
84    #[must_use]
85    pub fn with_policy(mut self, policy: TemporalPolicy) -> Self {
86        self.policy = policy;
87        self
88    }
89
90    /// Set target population.
91    #[must_use]
92    pub fn with_target_population(mut self, population: TargetPopulation) -> Self {
93        self.target_population = population;
94        self
95    }
96
97    /// Validate treatment/outcome, interventions, policy, and horizon.
98    ///
99    /// # Errors
100    ///
101    /// Returns [`QueryError`] on inconsistent configuration.
102    pub fn validate(&self) -> Result<(), QueryError> {
103        if self.treatment == self.outcome {
104            return Err(QueryError::TreatmentEqualsOutcome { id: self.treatment });
105        }
106        if self.horizon_steps == 0 {
107            return Err(QueryError::NonPositiveHorizon);
108        }
109        self.policy.validate().map_err(|e| match e {
110            crate::intervention::InterventionError::InvalidTemporalWindow { from, until } => {
111                QueryError::InvalidTemporalWindow { from, until }
112            }
113            other => QueryError::InvalidIntervention(other.to_string()),
114        })?;
115        self.target_population.validate()?;
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        Ok(())
133    }
134
135    /// Treatment time offset for Pulse `at` / Sustained `from` / Dynamic first active step.
136    #[must_use]
137    pub fn treatment_offset(&self) -> i32 {
138        self.try_treatment_offset().unwrap_or_default()
139    }
140
141    /// Treatment time offset when the policy has a defined origin.
142    ///
143    /// # Errors
144    ///
145    /// [`QueryError::DynamicPolicyHasNoTreatmentOffset`] for an empty dynamic schedule.
146    pub fn try_treatment_offset(&self) -> Result<i32, QueryError> {
147        match &self.policy {
148            TemporalPolicy::Pulse { at } => Ok(*at),
149            TemporalPolicy::Sustained { from, .. } => Ok(*from),
150            TemporalPolicy::Dynamic { active_at, .. } => {
151                active_at.first().copied().ok_or(QueryError::DynamicPolicyHasNoTreatmentOffset)
152            }
153        }
154    }
155
156    /// Outcome evaluation offset: `horizon_steps - 1` (absolute from window origin).
157    #[must_use]
158    pub fn outcome_offset(&self) -> i32 {
159        i32::try_from(self.horizon_steps.saturating_sub(1)).unwrap_or(i32::MAX)
160    }
161}
antecedent_core/query/temporal.rs

antecedent_core/query/
temporal.rs
