causal_hub/models/bayesian_network/mixed/
parameters.rs

1use std::borrow::Cow;
2
3use approx::{AbsDiffEq, RelativeEq};
4use rand::Rng;
5use serde::{Deserialize, Serialize};
6
7use crate::{
8    datasets::{CatSample, GaussSample},
9    models::{CPD, CatCPD, CatCPDS, CatSupport, GaussCPD, GaussCPDS, GaussSupport, HasLabels},
10    types::{Error, Labels, Result},
11};
12
13/// Unified support metadata for mixed CPDs.
14#[non_exhaustive]
15#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
16#[serde(untagged)]
17pub enum MixedSupport {
18    /// Categorical support (discrete states).
19    Categorical(CatSupport),
20    /// Gaussian support (continuous ranges).
21    Gaussian(GaussSupport),
22}
23
24/// The parameters of a mixed CPD.
25#[non_exhaustive]
26#[derive(Clone, Debug, Serialize, Deserialize)]
27#[serde(untagged)]
28pub enum MixedCPD {
29    /// Categorical CPD.
30    Categorical(CatCPD),
31    /// Gaussian CPD.
32    Gaussian(GaussCPD),
33}
34
35/// The sufficient statistics of a mixed CPD.
36#[non_exhaustive]
37#[derive(Clone, Debug, Serialize, Deserialize)]
38#[serde(untagged)]
39pub enum MixedCPDS {
40    /// Categorical sufficient statistics.
41    Categorical(CatCPDS),
42    /// Gaussian sufficient statistics.
43    Gaussian(Box<GaussCPDS>),
44}
45
46/// A unified sample type for mixed Bayesian networks.
47#[non_exhaustive]
48#[derive(Clone, Debug, Serialize, Deserialize)]
49#[serde(untagged)]
50pub enum MixedSample {
51    /// Categorical sample.
52    Categorical(CatSample),
53    /// Gaussian sample.
54    Gaussian(GaussSample),
55}
56
57impl From<CatCPD> for MixedCPD {
58    #[inline]
59    fn from(distribution: CatCPD) -> Self {
60        Self::Categorical(distribution)
61    }
62}
63
64impl From<GaussCPD> for MixedCPD {
65    #[inline]
66    fn from(distribution: GaussCPD) -> Self {
67        Self::Gaussian(distribution)
68    }
69}
70
71impl From<CatCPDS> for MixedCPDS {
72    #[inline]
73    fn from(stats: CatCPDS) -> Self {
74        Self::Categorical(stats)
75    }
76}
77
78impl From<GaussCPDS> for MixedCPDS {
79    #[inline]
80    fn from(stats: GaussCPDS) -> Self {
81        Self::Gaussian(Box::new(stats))
82    }
83}
84
85impl From<CatSample> for MixedSample {
86    #[inline]
87    fn from(sample: CatSample) -> Self {
88        Self::Categorical(sample)
89    }
90}
91
92impl From<GaussSample> for MixedSample {
93    #[inline]
94    fn from(sample: GaussSample) -> Self {
95        Self::Gaussian(sample)
96    }
97}
98
99impl HasLabels for MixedCPD {
100    fn labels(&self) -> &Labels {
101        match self {
102            Self::Categorical(distribution) => distribution.labels(),
103            Self::Gaussian(distribution) => distribution.labels(),
104        }
105    }
106}
107
108impl PartialEq for MixedCPD {
109    fn eq(&self, other: &Self) -> bool {
110        match (self, other) {
111            (Self::Categorical(a), Self::Categorical(b)) => a.eq(b),
112            (Self::Gaussian(a), Self::Gaussian(b)) => a.eq(b),
113            _ => false,
114        }
115    }
116}
117
118impl AbsDiffEq for MixedCPD {
119    type Epsilon = f64;
120
121    fn default_epsilon() -> Self::Epsilon {
122        Self::Epsilon::default_epsilon()
123    }
124
125    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
126        match (self, other) {
127            (Self::Categorical(a), Self::Categorical(b)) => a.abs_diff_eq(b, epsilon),
128            (Self::Gaussian(a), Self::Gaussian(b)) => a.abs_diff_eq(b, epsilon),
129            _ => false,
130        }
131    }
132}
133
134impl RelativeEq for MixedCPD {
135    fn default_max_relative() -> Self::Epsilon {
136        Self::Epsilon::default_max_relative()
137    }
138
139    fn relative_eq(
140        &self,
141        other: &Self,
142        epsilon: Self::Epsilon,
143        max_relative: Self::Epsilon,
144    ) -> bool {
145        match (self, other) {
146            (Self::Categorical(a), Self::Categorical(b)) => a.relative_eq(b, epsilon, max_relative),
147            (Self::Gaussian(a), Self::Gaussian(b)) => a.relative_eq(b, epsilon, max_relative),
148            _ => false,
149        }
150    }
151}
152
153impl CPD for MixedCPD {
154    type Sample = MixedSample;
155    type Support = MixedSupport;
156    type Parameters = MixedCPD;
157    type Statistics = MixedCPDS;
158
159    fn conditioning_labels(&self) -> &Labels {
160        match self {
161            Self::Categorical(distribution) => distribution.conditioning_labels(),
162            Self::Gaussian(distribution) => distribution.conditioning_labels(),
163        }
164    }
165
166    fn support(&self) -> Cow<'_, Self::Support> {
167        match self {
168            Self::Categorical(distribution) => {
169                Cow::Owned(MixedSupport::Categorical(distribution.support().clone()))
170            }
171            Self::Gaussian(distribution) => {
172                Cow::Owned(MixedSupport::Gaussian(distribution.support().into_owned()))
173            }
174        }
175    }
176
177    fn conditioning_support(&self) -> Cow<'_, Self::Support> {
178        match self {
179            Self::Categorical(distribution) => Cow::Owned(MixedSupport::Categorical(
180                distribution.conditioning_support().clone(),
181            )),
182            Self::Gaussian(distribution) => Cow::Owned(MixedSupport::Gaussian(
183                distribution.conditioning_support().into_owned(),
184            )),
185        }
186    }
187
188    fn parameters(&self) -> &Self::Parameters {
189        self
190    }
191
192    fn parameters_size(&self) -> usize {
193        match self {
194            Self::Categorical(distribution) => distribution.parameters_size(),
195            Self::Gaussian(distribution) => distribution.parameters_size(),
196        }
197    }
198
199    fn fitted_statistics(&self) -> Option<Cow<'_, Self::Statistics>> {
200        match self {
201            Self::Categorical(distribution) => distribution
202                .fitted_statistics()
203                .map(|stats| Cow::Owned(MixedCPDS::Categorical(stats.into_owned()))),
204            Self::Gaussian(distribution) => distribution
205                .fitted_statistics()
206                .map(|stats| Cow::Owned(MixedCPDS::Gaussian(Box::new(stats.into_owned())))),
207        }
208    }
209
210    fn fitted_log_likelihood(&self) -> Option<f64> {
211        match self {
212            Self::Categorical(distribution) => distribution.fitted_log_likelihood(),
213            Self::Gaussian(distribution) => distribution.fitted_log_likelihood(),
214        }
215    }
216
217    fn pf(&self, x: &Self::Sample, z: &Self::Sample) -> Result<f64> {
218        match (self, x, z) {
219            (
220                Self::Categorical(distribution),
221                MixedSample::Categorical(x),
222                MixedSample::Categorical(z),
223            ) => distribution.pf(x, z),
224            (Self::Gaussian(distribution), MixedSample::Gaussian(x), MixedSample::Gaussian(z)) => {
225                distribution.pf(x, z)
226            }
227            _ => Err(Error::InvalidParameter(
228                "x/z",
229                "sample type must match the CPD parameter type",
230            )),
231        }
232    }
233
234    fn sample<R: Rng>(&self, rng: &mut R, z: &Self::Sample) -> Result<Self::Sample> {
235        match (self, z) {
236            (Self::Categorical(distribution), MixedSample::Categorical(z)) => {
237                let sample = distribution.sample(rng, z)?;
238                Ok(MixedSample::Categorical(sample))
239            }
240            (Self::Gaussian(distribution), MixedSample::Gaussian(z)) => {
241                let sample = distribution.sample(rng, z)?;
242                Ok(MixedSample::Gaussian(sample))
243            }
244            _ => Err(Error::InvalidParameter(
245                "z",
246                "sample type must match the CPD parameter type",
247            )),
248        }
249    }
250}
causal_hub/models/bayesian_network/mixed/parameters.rs

causal_hub/models/bayesian_network/mixed/
parameters.rs
