causal_hub/estimators/parameters/mod.rs
1//! Parameter estimation: maximum-likelihood, Bayesian, sufficient-statistics, EM, and RAWE estimators.
2
3mod bayesian;
4pub use bayesian::*;
5
6mod expectation_maximization;
7pub use expectation_maximization::*;
8
9mod maximum_likelihood;
10pub use maximum_likelihood::*;
11
12mod sufficient_statistics;
13pub use sufficient_statistics::*;
14
15mod raw;
16pub use raw::*;
17use rayon::prelude::*;
18
19use crate::{
20 models::{BN, CIM, CPD, CTBN, DiGraph, Graph},
21 set,
22 types::{Result, Set},
23};
24
25/// A trait for sufficient statistics estimators.
26pub trait CSSEstimator<T> {
27 /// Fits the estimator to the dataset and returns the conditional sufficient statistics.
28 ///
29 /// # Arguments
30 ///
31 /// * `x` - The variable to fit the estimator to.
32 /// * `z` - The variables to condition on.
33 ///
34 /// # Errors
35 ///
36 /// * If the set of variables to fit the estimator to is empty.
37 /// * If the set of variables to fit the estimator to is not a subset of the dataset variables.
38 /// * If the set of variables to condition on is not a subset of the dataset variables.
39 /// * If the sets of variables are not disjoint.
40 ///
41 /// # Returns
42 ///
43 /// The sufficient statistics.
44 ///
45 fn fit(&self, x: &Set<usize>, z: &Set<usize>) -> Result<T>;
46}
47
48/// A trait for sufficient statistics estimators in parallel.
49pub trait ParCSSEstimator<T> {
50 /// Fits the estimator to the dataset and returns the conditional sufficient statistics in parallel.
51 ///
52 /// # Arguments
53 ///
54 /// * `x` - The variable to fit the estimator to.
55 /// * `z` - The variables to condition on.
56 ///
57 /// # Errors
58 ///
59 /// * If the set of variables to fit the estimator to is empty.
60 /// * If the set of variables to fit the estimator to is not a subset of the dataset variables.
61 /// * If the set of variables to condition on is not a subset of the dataset variables.
62 /// * If the sets of variables are not disjoint.
63 ///
64 /// # Returns
65 ///
66 /// The sufficient statistics.
67 ///
68 fn par_fit(&self, x: &Set<usize>, z: &Set<usize>) -> Result<T>;
69}
70
71/// A trait for conditional probability distribution estimators.
72pub trait CPDEstimator<T> {
73 /// Fits the estimator to the dataset and returns a CPD.
74 ///
75 /// # Arguments
76 ///
77 /// * `x` - The variable to fit the estimator to.
78 /// * `z` - The variables to condition on.
79 ///
80 /// # Errors
81 ///
82 /// * If the set of variables to fit the estimator to is empty.
83 /// * If the set of variables to fit the estimator to is not a subset of the dataset variables.
84 /// * If the set of variables to condition on is not a subset of the dataset variables.
85 /// * If the sets of variables are not disjoint.
86 ///
87 /// # Returns
88 ///
89 /// The estimated CPD.
90 ///
91 fn fit(&self, x: &Set<usize>, z: &Set<usize>) -> Result<T>;
92}
93
94/// A trait for conditional probability distribution estimators in parallel.
95pub trait ParCPDEstimator<T> {
96 /// Fits the estimator to the dataset and returns a CPD in parallel.
97 ///
98 /// # Arguments
99 ///
100 /// * `x` - The variable to fit the estimator to.
101 /// * `z` - The variables to condition on.
102 ///
103 /// # Errors
104 ///
105 /// * If the set of variables to fit the estimator to is empty.
106 /// * If the set of variables to fit the estimator to is not a subset of the dataset variables.
107 /// * If the set of variables to condition on is not a subset of the dataset variables.
108 /// * If the sets of variables are not disjoint.
109 ///
110 /// # Returns
111 ///
112 /// The estimated CPD.
113 ///
114 fn par_fit(&self, x: &Set<usize>, z: &Set<usize>) -> Result<T>;
115}
116
117/// A trait for Bayesian network estimators.
118pub trait BNEstimator<T> {
119 /// Fits the estimator to the dataset and returns a Bayesian network.
120 ///
121 /// # Arguments
122 ///
123 /// * `graph` - The graph to fit the estimator to.
124 ///
125 /// # Errors
126 ///
127 /// * If the graph labels are not a subset of the dataset variables.
128 ///
129 /// # Returns
130 ///
131 /// The estimated Bayesian network.
132 ///
133 fn fit(&self, graph: DiGraph) -> Result<T>;
134}
135
136/// Blanket implement for all BN estimators with a corresponding CPD estimator.
137impl<T, E> BNEstimator<T> for E
138where
139 T: BN,
140 T::CPD: CPD,
141 E: CPDEstimator<T::CPD>,
142{
143 fn fit(&self, graph: DiGraph) -> Result<T> {
144 // Fit the parameters of the distribution using the estimator.
145 let cpds: Vec<_> = graph
146 .vertices()
147 .into_iter()
148 .map(|i| {
149 let i = set![i];
150 self.fit(&i, &graph.parents(&i)?)
151 })
152 .collect::<Result<_>>()?;
153 // Construct the BN with the graph and the parameters.
154 T::new(graph, cpds)
155 }
156}
157
158/// A trait for parallel Bayesian network estimators.
159pub trait ParBNEstimator<T> {
160 /// Fits the estimator to the dataset and returns a Bayesian network in parallel.
161 ///
162 /// # Arguments
163 ///
164 /// * `graph` - The graph to fit the estimator to.
165 ///
166 /// # Errors
167 ///
168 /// * If the graph labels are not a subset of the dataset variables.
169 ///
170 /// # Returns
171 ///
172 /// The estimated Bayesian network.
173 ///
174 fn par_fit(&self, graph: DiGraph) -> Result<T>;
175}
176
177/// Blanket implement for all BN estimators with a corresponding CPD estimator.
178impl<T, E> ParBNEstimator<T> for E
179where
180 T: BN,
181 T::CPD: CPD + Send,
182 E: ParCPDEstimator<T::CPD> + Sync,
183{
184 fn par_fit(&self, graph: DiGraph) -> Result<T> {
185 // Fit the parameters of the distribution using the estimator.
186 let cpds: Vec<_> = graph
187 .vertices()
188 .into_par_iter()
189 .map(|i| {
190 let i = set![i];
191 self.par_fit(&i, &graph.parents(&i)?)
192 })
193 .collect::<Result<_>>()?;
194 // Construct the BN with the graph and the parameters.
195 T::new(graph, cpds)
196 }
197}
198
199/// A trait for CTBN estimators.
200pub trait CTBNEstimator<T> {
201 /// Fits the estimator to the trajectory and returns a CTBN.
202 ///
203 /// # Arguments
204 ///
205 /// * `graph` - The graph to fit the estimator to.
206 ///
207 /// # Errors
208 ///
209 /// * If the graph labels are not a subset of the dataset variables.
210 ///
211 /// # Returns
212 ///
213 /// The estimated CTBN.
214 ///
215 fn fit(&self, graph: DiGraph) -> Result<T>;
216}
217
218/// Blanket implement for all CTBN estimators with a corresponding CIM estimator.
219impl<T, E> CTBNEstimator<T> for E
220where
221 T: CTBN,
222 T::CIM: CIM,
223 E: CPDEstimator<T::CIM>,
224{
225 fn fit(&self, graph: DiGraph) -> Result<T> {
226 // Fit the parameters of the distribution using the estimator.
227 let cims: Vec<_> = graph
228 .vertices()
229 .into_iter()
230 .map(|i| {
231 let i = set![i];
232 self.fit(&i, &graph.parents(&i)?)
233 })
234 .collect::<Result<_>>()?;
235 // Construct the CTBN with the graph and the parameters.
236 T::new(graph, cims)
237 }
238}
239
240/// A trait for parallel CTBN estimators.
241pub trait ParCTBNEstimator<T> {
242 /// Fits the estimator to the trajectory and returns a CTBN in parallel.
243 ///
244 /// # Arguments
245 ///
246 /// * `graph` - The graph to fit the estimator to.
247 ///
248 /// # Errors
249 ///
250 /// * If the graph labels are not a subset of the dataset variables.
251 ///
252 /// # Returns
253 ///
254 /// The estimated CTBN.
255 ///
256 fn par_fit(&self, graph: DiGraph) -> Result<T>;
257}
258
259/// Blanket implement for all CTBN estimators with a corresponding CIM estimator.
260impl<T, E> ParCTBNEstimator<T> for E
261where
262 T: CTBN,
263 T::CIM: CIM + Send,
264 E: ParCPDEstimator<T::CIM> + Sync,
265{
266 fn par_fit(&self, graph: DiGraph) -> Result<T> {
267 // Fit the parameters of the distribution using the estimator.
268 let cims: Vec<_> = graph
269 .vertices()
270 .into_par_iter()
271 .map(|i| {
272 let i = set![i];
273 self.par_fit(&i, &graph.parents(&i)?)
274 })
275 .collect::<Result<_>>()?;
276 // Construct the CTBN with the graph and the parameters.
277 T::new(graph, cims)
278 }
279}