Skip to main content

causal_hub/estimators/parameters/bayesian/
mod.rs

1//! Bayesian (categorical / gaussian / mixed) parameter estimators.
2
3mod table;
4mod trajectory;
5
6use crate::{
7    datasets::{MissingMechanism, MissingMethod},
8    models::HasLabels,
9    types::{Error, Labels, Result},
10};
11
12/// A struct representing a Bayesian estimator.
13#[derive(Clone, Debug)]
14pub struct BE<'a, D, T> {
15    dataset: &'a D,
16    missing_method: Option<MissingMethod>,
17    missing_mechanism: Option<MissingMechanism>,
18    prior: T,
19}
20
21impl<'a, D> BE<'a, D, ()> {
22    /// Creates a new Bayesian estimator.
23    ///
24    /// # Arguments
25    ///
26    /// * `dataset` - A reference to the dataset to fit the estimator to.
27    ///
28    /// # Returns
29    ///
30    /// A new Bayesian estimator.
31    ///
32    #[inline]
33    pub const fn new(dataset: &'a D) -> Self {
34        Self {
35            dataset,
36            missing_method: None,
37            missing_mechanism: None,
38            prior: (),
39        }
40    }
41}
42
43impl<'a, D, T> BE<'a, D, T> {
44    /// Sets the missing handling method.
45    ///
46    /// # Arguments
47    ///
48    /// * `missing_method` - An optional missing handling method to set.
49    /// * `missing_mechanism` - An optional missing mechanism to set.
50    ///
51    /// # Returns
52    ///
53    /// A new estimator with the specified missing handling method.
54    ///
55    #[inline]
56    pub fn with_missing_method(
57        mut self,
58        missing_method: Option<MissingMethod>,
59        missing_mechanism: Option<MissingMechanism>,
60    ) -> Result<Self> {
61        // Validate missing method and mechanism.
62        match (missing_method, &missing_mechanism) {
63            (Some(MissingMethod::LW) | Some(MissingMethod::PW), Some(_)) => {
64                return Err(Error::InvalidParameter(
65                    "missing_mechanism",
66                    "must be None if missing_method is LW or PW",
67                ));
68            }
69            (Some(MissingMethod::IPW) | Some(MissingMethod::AIPW), None) => {
70                return Err(Error::InvalidParameter(
71                    "missing_mechanism",
72                    "must be provided if missing_method is IPW or AIPW",
73                ));
74            }
75            _ => {}
76        }
77
78        self.missing_method = missing_method;
79        self.missing_mechanism = missing_mechanism;
80        Ok(self)
81    }
82
83    /// Sets the prior distribution.
84    ///
85    /// # Arguments
86    ///
87    /// * `prior` - The prior distribution to set.
88    ///
89    /// # Returns
90    ///
91    /// A new Bayesian estimator with the specified prior.
92    ///
93    #[inline]
94    pub fn with_prior<U>(self, prior: U) -> BE<'a, D, U> {
95        BE {
96            dataset: self.dataset,
97            missing_method: self.missing_method,
98            missing_mechanism: self.missing_mechanism,
99            prior,
100        }
101    }
102}
103
104impl<D, T> HasLabels for BE<'_, D, T>
105where
106    D: HasLabels,
107{
108    #[inline]
109    fn labels(&self) -> &Labels {
110        self.dataset.labels()
111    }
112}