antecedent_core/query/
population.rs

1//! Target-population resolution against tabular rows.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use std::collections::BTreeMap;
6use std::sync::Arc;
7
8use crate::ids::DistributionRef;
9
10use super::error::QueryError;
11use super::target::{PredicateExpr, TargetPopulation};
12
13/// Caller-supplied bindings for named predicates and custom target distributions.
14#[derive(Clone, Debug, Default)]
15pub struct PopulationRegistry {
16    predicates: BTreeMap<Arc<str>, Arc<[usize]>>,
17    distributions: BTreeMap<u32, Arc<[f64]>>,
18}
19
20impl PopulationRegistry {
21    /// Empty registry.
22    #[must_use]
23    pub fn new() -> Self {
24        Self::default()
25    }
26
27    /// Bind a named predicate to explicit row indexes.
28    pub fn insert_predicate(&mut self, name: impl Into<Arc<str>>, rows: impl Into<Arc<[usize]>>) {
29        self.predicates.insert(name.into(), rows.into());
30    }
31
32    /// Bind a custom distribution handle to non-negative row weights (length = `n`).
33    pub fn insert_distribution(&mut self, id: DistributionRef, weights: impl Into<Arc<[f64]>>) {
34        self.distributions.insert(id.raw(), weights.into());
35    }
36
37    /// Look up a named predicate.
38    #[must_use]
39    pub fn predicate(&self, name: &str) -> Option<&[usize]> {
40        self.predicates.get(name).map(std::convert::AsRef::as_ref)
41    }
42
43    /// Look up distribution weights.
44    #[must_use]
45    pub fn distribution(&self, id: DistributionRef) -> Option<&[f64]> {
46        self.distributions.get(&id.raw()).map(std::convert::AsRef::as_ref)
47    }
48}
49
50/// Resolved population as a keep-mask and optional observation weights.
51#[derive(Clone, Debug)]
52pub struct PopulationSelection {
53    /// Length-`n` keep mask (`true` = include).
54    pub keep: Arc<[bool]>,
55    /// Optional length-`n` non-negative weights (`CustomDistribution`).
56    pub weights: Option<Arc<[f64]>>,
57}
58
59impl TargetPopulation {
60    /// Resolve this population over `n` rows.
61    ///
62    /// `treatment` is required for [`Self::Treated`] / [`Self::Untreated`] (binary 0/1 column,
63    /// length `n`). Named predicates and custom distributions require `registry`.
64    ///
65    /// # Errors
66    ///
67    /// Missing registry bindings, shape mismatches, or unknown variants.
68    pub fn resolve(
69        &self,
70        n: usize,
71        treatment: Option<&[f64]>,
72        registry: Option<&PopulationRegistry>,
73    ) -> Result<PopulationSelection, QueryError> {
74        match self {
75            Self::AllObserved => {
76                Ok(PopulationSelection { keep: Arc::from(vec![true; n]), weights: None })
77            }
78            Self::Treated | Self::Untreated => {
79                let t = treatment.ok_or(QueryError::PopulationNeedsTreatment)?;
80                if t.len() != n {
81                    return Err(QueryError::PopulationLengthMismatch {
82                        expected: n,
83                        actual: t.len(),
84                    });
85                }
86                let want_treated = matches!(self, Self::Treated);
87                let mut keep = vec![false; n];
88                for (i, &ti) in t.iter().enumerate() {
89                    let is_t = (ti - 1.0).abs() <= 1e-12;
90                    let is_c = ti.abs() <= 1e-12;
91                    if !is_t && !is_c {
92                        return Err(QueryError::PopulationNonBinaryTreatment);
93                    }
94                    keep[i] = if want_treated { is_t } else { is_c };
95                }
96                Ok(PopulationSelection { keep: Arc::from(keep), weights: None })
97            }
98            Self::Environment(_) => Err(QueryError::PopulationEnvironmentUnsupported),
99            Self::Predicate(expr) => {
100                let rows: &[usize] = match expr {
101                    PredicateExpr::Rows(rows) => rows.as_ref(),
102                    PredicateExpr::Named(name) => {
103                        let reg = registry.ok_or(QueryError::PopulationRegistryRequired)?;
104                        reg.predicate(name.as_ref()).ok_or_else(|| {
105                            QueryError::UnknownPredicateName { name: Arc::clone(name) }
106                        })?
107                    }
108                };
109                let mut keep = vec![false; n];
110                for &r in rows {
111                    if r >= n {
112                        return Err(QueryError::PopulationRowOutOfRange { row: r, n });
113                    }
114                    keep[r] = true;
115                }
116                Ok(PopulationSelection { keep: Arc::from(keep), weights: None })
117            }
118            Self::CustomDistribution(id) => {
119                let reg = registry.ok_or(QueryError::PopulationRegistryRequired)?;
120                let weights = reg
121                    .distribution(*id)
122                    .ok_or(QueryError::UnknownDistributionRef { id: id.raw() })?;
123                if weights.len() != n {
124                    return Err(QueryError::PopulationLengthMismatch {
125                        expected: n,
126                        actual: weights.len(),
127                    });
128                }
129                if weights.iter().any(|w| !w.is_finite() || *w < 0.0) {
130                    return Err(QueryError::InvalidPopulationWeights);
131                }
132                Ok(PopulationSelection {
133                    keep: Arc::from(vec![true; n]),
134                    weights: Some(Arc::from(weights.to_vec())),
135                })
136            }
137        }
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::ids::DistributionRef;
145
146    #[test]
147    fn resolves_rows_and_named_and_weights() {
148        let mut reg = PopulationRegistry::new();
149        reg.insert_predicate("cohort", [0usize, 2]);
150        reg.insert_distribution(DistributionRef::from_raw(1), [0.5, 0.0, 1.5]);
151
152        let rows = TargetPopulation::Predicate(PredicateExpr::rows([1usize]))
153            .resolve(3, None, None)
154            .unwrap();
155        assert_eq!(rows.keep.as_ref(), &[false, true, false]);
156
157        let named = TargetPopulation::Predicate(PredicateExpr::named("cohort"))
158            .resolve(3, None, Some(&reg))
159            .unwrap();
160        assert_eq!(named.keep.as_ref(), &[true, false, true]);
161
162        let w = TargetPopulation::CustomDistribution(DistributionRef::from_raw(1))
163            .resolve(3, None, Some(&reg))
164            .unwrap();
165        assert_eq!(w.weights.as_ref().unwrap().as_ref(), &[0.5, 0.0, 1.5]);
166    }
167}
antecedent_core/query/population.rs

antecedent_core/query/
population.rs
