Skip to main content

causal_hub/estimators/parameters/maximum_likelihood/
mod.rs

1//! Maximum-likelihood 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 maximum likelihood estimator.
13#[derive(Clone, Debug)]
14pub struct MLE<'a, D> {
15    dataset: &'a D,
16    missing_method: Option<MissingMethod>,
17    missing_mechanism: Option<MissingMechanism>,
18}
19
20impl<'a, D> MLE<'a, D> {
21    /// Creates a new maximum likelihood estimator.
22    ///
23    /// # Arguments
24    ///
25    /// * `dataset` - A reference to the dataset to fit the estimator to.
26    ///
27    /// # Returns
28    ///
29    /// A new `MaximumLikelihoodEstimator` instance.
30    ///
31    #[inline]
32    pub const fn new(dataset: &'a D) -> Self {
33        Self {
34            dataset,
35            missing_method: None,
36            missing_mechanism: None,
37        }
38    }
39
40    /// Sets the missing handling method.
41    ///
42    /// # Arguments
43    ///
44    /// * `missing_method` - An optional missing handling method to set.
45    /// * `missing_mechanism` - An optional missing mechanism to set.
46    ///
47    /// # Returns
48    ///
49    /// A new estimator with the specified missing handling method.
50    ///
51    #[inline]
52    pub fn with_missing_method(
53        mut self,
54        missing_method: Option<MissingMethod>,
55        missing_mechanism: Option<MissingMechanism>,
56    ) -> Result<Self> {
57        // Validate missing method and mechanism.
58        match (missing_method, &missing_mechanism) {
59            (Some(MissingMethod::LW) | Some(MissingMethod::PW), Some(_)) => {
60                return Err(Error::InvalidParameter(
61                    "missing_mechanism",
62                    "must be None if missing_method is LW or PW",
63                ));
64            }
65            (Some(MissingMethod::IPW) | Some(MissingMethod::AIPW), None) => {
66                return Err(Error::InvalidParameter(
67                    "missing_mechanism",
68                    "must be provided if missing_method is IPW or AIPW",
69                ));
70            }
71            _ => {}
72        }
73
74        self.missing_method = missing_method;
75        self.missing_mechanism = missing_mechanism;
76        Ok(self)
77    }
78}
79
80impl<D> HasLabels for MLE<'_, D>
81where
82    D: HasLabels,
83{
84    #[inline]
85    fn labels(&self) -> &Labels {
86        self.dataset.labels()
87    }
88}