causal_hub/models/bayesian_network/mixed/
potential.rs

1use std::{
2    borrow::Cow,
3    ops::{Div, DivAssign, Mul, MulAssign},
4};
5
6use approx::{AbsDiffEq, RelativeEq};
7use serde::{Deserialize, Serialize};
8
9use crate::{
10    impl_json_io,
11    models::{CatPhi, GaussPhi, HasLabels, Phi},
12    types::{Error, Labels, Result, Set},
13};
14
15/// A unified potential for mixed Bayesian networks.
16#[non_exhaustive]
17#[derive(Clone, Debug, Serialize, Deserialize)]
18#[serde(untagged)]
19pub enum MixedPhi {
20    /// Categorical potential.
21    Categorical(CatPhi),
22    /// Gaussian potential.
23    Gaussian(GaussPhi),
24}
25
26impl HasLabels for MixedPhi {
27    #[inline]
28    fn labels(&self) -> &Labels {
29        match self {
30            Self::Categorical(potential) => potential.labels(),
31            Self::Gaussian(potential) => potential.labels(),
32        }
33    }
34}
35
36impl PartialEq for MixedPhi {
37    fn eq(&self, other: &Self) -> bool {
38        match (self, other) {
39            (Self::Categorical(a), Self::Categorical(b)) => a.eq(b),
40            (Self::Gaussian(a), Self::Gaussian(b)) => a.eq(b),
41            _ => false,
42        }
43    }
44}
45
46impl AbsDiffEq for MixedPhi {
47    type Epsilon = f64;
48
49    fn default_epsilon() -> Self::Epsilon {
50        Self::Epsilon::default_epsilon()
51    }
52
53    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
54        match (self, other) {
55            (Self::Categorical(a), Self::Categorical(b)) => a.abs_diff_eq(b, epsilon),
56            (Self::Gaussian(a), Self::Gaussian(b)) => a.abs_diff_eq(b, epsilon),
57            _ => false,
58        }
59    }
60}
61
62impl RelativeEq for MixedPhi {
63    fn default_max_relative() -> Self::Epsilon {
64        Self::Epsilon::default_max_relative()
65    }
66
67    fn relative_eq(
68        &self,
69        other: &Self,
70        epsilon: Self::Epsilon,
71        max_relative: Self::Epsilon,
72    ) -> bool {
73        match (self, other) {
74            (Self::Categorical(a), Self::Categorical(b)) => a.relative_eq(b, epsilon, max_relative),
75            (Self::Gaussian(a), Self::Gaussian(b)) => a.relative_eq(b, epsilon, max_relative),
76            _ => false,
77        }
78    }
79}
80
81impl Phi for MixedPhi {
82    type CPD = crate::models::MixedCPD;
83    type Support = crate::models::MixedSupport;
84    type Parameters = MixedPhi;
85    type Evidence = crate::models::MixedEv;
86
87    fn support(&self) -> Cow<'_, Self::Support> {
88        match self {
89            Self::Categorical(potential) => Cow::Owned(crate::models::MixedSupport::Categorical(
90                potential.support().clone(),
91            )),
92            Self::Gaussian(potential) => Cow::Owned(crate::models::MixedSupport::Gaussian(
93                Cow::into_owned(potential.support()),
94            )),
95        }
96    }
97
98    fn parameters(&self) -> &Self::Parameters {
99        self
100    }
101
102    fn parameters_size(&self) -> usize {
103        match self {
104            Self::Categorical(potential) => potential.parameters_size(),
105            Self::Gaussian(potential) => potential.parameters_size(),
106        }
107    }
108
109    fn condition(&self, evidence: &Self::Evidence) -> Result<Self> {
110        match (self, evidence) {
111            (Self::Categorical(potential), crate::models::MixedEv::Categorical(ev)) => {
112                potential.condition(ev).map(Self::Categorical)
113            }
114            (Self::Gaussian(potential), crate::models::MixedEv::Gaussian(ev)) => {
115                potential.condition(ev).map(Self::Gaussian)
116            }
117            _ => Err(Error::InvalidParameter(
118                "e",
119                "evidence type must match the potential variant",
120            )),
121        }
122    }
123
124    fn marginalize(&self, x: &Set<usize>) -> Result<Self> {
125        match self {
126            Self::Categorical(potential) => potential.marginalize(x).map(Self::Categorical),
127            Self::Gaussian(potential) => potential.marginalize(x).map(Self::Gaussian),
128        }
129    }
130
131    fn normalize(&self) -> Result<Self> {
132        match self {
133            Self::Categorical(potential) => potential.normalize().map(Self::Categorical),
134            Self::Gaussian(potential) => potential.normalize().map(Self::Gaussian),
135        }
136    }
137
138    fn from_cpd(distribution: Self::CPD) -> Result<Self> {
139        match distribution {
140            crate::models::MixedCPD::Categorical(c) => c.into_phi().map(Self::Categorical),
141            crate::models::MixedCPD::Gaussian(c) => c.into_phi().map(Self::Gaussian),
142        }
143    }
144
145    fn into_cpd(self, x: &Set<usize>, z: &Set<usize>) -> Result<Self::CPD> {
146        match self {
147            Self::Categorical(potential) => potential
148                .into_cpd(x, z)
149                .map(crate::models::MixedCPD::Categorical),
150            Self::Gaussian(potential) => potential
151                .into_cpd(x, z)
152                .map(crate::models::MixedCPD::Gaussian),
153        }
154    }
155}
156
157impl MulAssign<&MixedPhi> for MixedPhi {
158    fn mul_assign(&mut self, rhs: &MixedPhi) {
159        match (self, rhs) {
160            (Self::Categorical(a), MixedPhi::Categorical(b)) => {
161                a.mul_assign(b);
162            }
163            (Self::Gaussian(a), MixedPhi::Gaussian(b)) => {
164                a.mul_assign(b);
165            }
166            _ => unreachable!("cannot multiply mixed potential variants"),
167        }
168    }
169}
170
171impl Mul<&MixedPhi> for &MixedPhi {
172    type Output = MixedPhi;
173
174    #[inline]
175    fn mul(self, rhs: &MixedPhi) -> Self::Output {
176        let mut lhs = self.clone();
177        MulAssign::mul_assign(&mut lhs, rhs);
178        lhs
179    }
180}
181
182impl DivAssign<&MixedPhi> for MixedPhi {
183    fn div_assign(&mut self, rhs: &MixedPhi) {
184        match (self, rhs) {
185            (Self::Categorical(a), MixedPhi::Categorical(b)) => {
186                DivAssign::div_assign(a, b);
187            }
188            (Self::Gaussian(a), MixedPhi::Gaussian(b)) => {
189                DivAssign::div_assign(a, b);
190            }
191            _ => unreachable!("cannot divide mixed potential variants"),
192        }
193    }
194}
195
196impl Div<&MixedPhi> for &MixedPhi {
197    type Output = MixedPhi;
198
199    #[inline]
200    fn div(self, rhs: &MixedPhi) -> Self::Output {
201        let mut lhs = self.clone();
202        DivAssign::div_assign(&mut lhs, rhs);
203        lhs
204    }
205}
206
207// Implement `JsonIO` for `MixedPhi`.
208impl_json_io!(MixedPhi);
causal_hub/models/bayesian_network/mixed/potential.rs

causal_hub/models/bayesian_network/mixed/
potential.rs
