use dry::macro_for;
use ndarray::prelude::*;
use crate::{
datasets::{CatIncTable, CatTable, CatWtdTable},
estimators::{CPDEstimator, CSSEstimator, MLE, ParCPDEstimator, ParCSSEstimator, SSE},
models::{CatCPD, CatCPDS},
types::{Set, States},
};
impl MLE<'_, CatTable> {
fn fit(states: &States, x: &Set<usize>, z: &Set<usize>, sample_statistics: CatCPDS) -> CatCPD {
let n_xz = sample_statistics.sample_conditional_counts();
let n_z = &n_xz.sum_axis(Axis(1)).insert_axis(Axis(1));
assert!(
n_z.iter().all(|&x| x > 0.),
"Failed to get non-zero counts.",
);
let parameters = n_xz / n_z;
let eps = f64::MIN_POSITIVE;
let sample_log_likelihood = (n_xz * (¶meters + eps).ln()).sum();
let conditioning_states = z
.iter()
.map(|&i| {
let (k, v) = states.get_index(i).unwrap();
(k.clone(), v.clone())
})
.collect();
let states = x
.iter()
.map(|&i| {
let (k, v) = states.get_index(i).unwrap();
(k.clone(), v.clone())
})
.collect();
let sample_statistics = Some(sample_statistics);
let sample_log_likelihood = Some(sample_log_likelihood);
CatCPD::with_optionals(
states,
conditioning_states,
parameters,
sample_statistics,
sample_log_likelihood,
)
}
}
macro_for!($type in [CatTable, CatIncTable, CatWtdTable] {
impl CPDEstimator<CatCPD> for MLE<'_, $type> {
fn fit(&self, x: &Set<usize>, z: &Set<usize>) -> CatCPD {
let states = self.dataset.states();
let sample_statistics = SSE::new(self.dataset);
let sample_statistics = sample_statistics.with_missing_method(
self.missing_method,
self.missing_mechanism.clone()
);
let sample_statistics = sample_statistics.fit(x, z);
MLE::<'_, CatTable>::fit(states, x, z, sample_statistics)
}
}
impl ParCPDEstimator<CatCPD> for MLE<'_, $type> {
fn par_fit(&self, x: &Set<usize>, z: &Set<usize>) -> CatCPD {
let states = self.dataset.states();
let sample_statistics = SSE::new(self.dataset);
let sample_statistics = sample_statistics.with_missing_method(
self.missing_method,
self.missing_mechanism.clone()
);
let sample_statistics = sample_statistics.par_fit(x, z);
MLE::<'_, CatTable>::fit(states, x, z, sample_statistics)
}
}
});