use std::sync::Arc;
use crate::ids::VariableId;
use crate::intervention::Intervention;
use crate::value::Value;
use super::TargetPopulation;
use super::error::QueryError;
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
#[non_exhaustive]
pub struct AverageEffectQuery {
pub treatment: VariableId,
pub outcome: VariableId,
pub effect_modifiers: Arc<[VariableId]>,
pub control: Intervention,
pub active: Intervention,
pub target_population: TargetPopulation,
}
impl AverageEffectQuery {
#[must_use]
pub fn new(
treatment: VariableId,
outcome: VariableId,
effect_modifiers: impl Into<Arc<[VariableId]>>,
control: Intervention,
active: Intervention,
target_population: TargetPopulation,
) -> Self {
Self {
treatment,
outcome,
effect_modifiers: effect_modifiers.into(),
control,
active,
target_population,
}
}
#[must_use]
pub fn binary_ate(treatment: VariableId, outcome: VariableId) -> Self {
Self::new(
treatment,
outcome,
Arc::from([]),
Intervention::set(treatment, Value::f64(0.0)),
Intervention::set(treatment, Value::f64(1.0)),
TargetPopulation::AllObserved,
)
}
#[must_use]
pub fn with_levels(
treatment: VariableId,
outcome: VariableId,
control_level: f64,
active_level: f64,
) -> Self {
Self {
treatment,
outcome,
effect_modifiers: Arc::from([]),
control: Intervention::set(treatment, Value::f64(control_level)),
active: Intervention::set(treatment, Value::f64(active_level)),
target_population: TargetPopulation::AllObserved,
}
}
#[must_use]
pub fn with_effect_modifiers(mut self, modifiers: impl Into<Arc<[VariableId]>>) -> Self {
self.effect_modifiers = modifiers.into();
self
}
#[must_use]
pub fn with_target_population(mut self, population: TargetPopulation) -> Self {
self.target_population = population;
self
}
pub fn validate(&self) -> Result<(), QueryError> {
if self.treatment == self.outcome {
return Err(QueryError::TreatmentEqualsOutcome { id: self.treatment });
}
let control_var =
self.control.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
if control_var != self.treatment {
return Err(QueryError::InterventionVariableMismatch {
expected: self.treatment,
got: control_var,
});
}
let active_var =
self.active.primary_variable().ok_or(QueryError::AmbiguousInterventionTarget)?;
if active_var != self.treatment {
return Err(QueryError::InterventionVariableMismatch {
expected: self.treatment,
got: active_var,
});
}
if self.effect_modifiers.iter().any(|m| *m == self.treatment || *m == self.outcome) {
return Err(QueryError::ModifierOverlapsTreatmentOrOutcome);
}
self.target_population.validate()?;
Ok(())
}
}