logaritmisk/trueskill-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Rustify some code.

Browse Source
This commit is contained in:
Anders Olsson
2018-10-24 20:15:29 +02:00
parent 77523e9db4
commit bcdabf9fbb
3 changed files with 85 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

166
src/factor_graph.rs
View File

@@ -52,9 +52,7 @@ impl Variable {
pub fn update_value(&mut self, factor: usize, value: Gaussian) {
let old = self.factors[&factor];
let intermediate = value * old;
self.factors.insert(factor, intermediate / self.value);
self.factors.insert(factor, value * old / self.value);
self.value = value;
}
@@ -65,12 +63,8 @@ impl Variable {
pub fn update_message(&mut self, factor: usize, message: Gaussian) {
let old = self.factors[&factor];
let intermediate = self.value / old;
let value = intermediate * message;
self.value = value;
self.factors.insert(factor, message);
self.value = self.value / old * message;
}
pub fn get_message(&self, factor: usize) -> Gaussian {
@@ -135,18 +129,15 @@ impl LikelihoodFactor {
}
pub fn update_mean(&self, variable_arena: &mut VariableArena) {
let x = variable_arena
let (x, fx) = variable_arena
.get(self.value)
.map(|variable| variable.get_value())
.unwrap();
let fx = variable_arena
.get_mut(self.value)
.map(|variable| variable.get_message(self.id))
.map(|variable| (
variable.get_value(),
variable.get_message(self.id)
))
.unwrap();
let a = 1.0 / (1.0 + self.variance * (x.pi() - fx.pi()));
let gaussian = Gaussian::from_pi_tau(a * (x.pi() - fx.pi()), a * (x.tau() - fx.tau()));
variable_arena
@@ -156,18 +147,15 @@ impl LikelihoodFactor {
}
pub fn update_value(&self, variable_arena: &mut VariableArena) {
let y = variable_arena
let (y, fy) = variable_arena
.get(self.mean)
.map(|variable| variable.get_value())
.unwrap();
let fy = variable_arena
.get(self.mean)
.map(|variable| variable.get_message(self.id))
.map(|variable| (
variable.get_value(),
variable.get_message(self.id)
))
.unwrap();
let a = 1.0 / (1.0 + self.variance * (y.pi() - fy.pi()));
let gaussian = Gaussian::from_pi_tau(a * (y.pi() - fy.pi()), a * (y.tau() - fy.tau()));
variable_arena
@@ -212,21 +200,17 @@ impl SumFactor {
variable: VariableId,
y: Vec<Gaussian>,
fy: Vec<Gaussian>,
a: &Vec<f64>,
a: &[f64],
) {
let size = a.len();
let (sum_pi, sum_tau) = a.iter().zip(y.iter().zip(fy.iter()))
.fold((0.0, 0.0), |(pi, tau), (a, (y, fy))| {
let x = *y / *fy;
let mut sum_pi = 0.0;
let mut sum_tau = 0.0;
for i in 0..size {
let da = a[i];
let gy = y[i];
let gfy = fy[i];
sum_pi += da.powi(2) / (gy.pi() - gfy.pi());
sum_tau += da * (gy.tau() - gfy.tau()) / (gy.pi() - gfy.pi());
}
(
pi + a.powi(2) / x.pi(),
tau + a * x.tau() / x.pi()
)
});
let new_pi = 1.0 / sum_pi;
let new_tau = new_pi * sum_tau;
@@ -240,69 +224,57 @@ impl SumFactor {
}
pub fn update_sum(&self, variable_arena: &mut VariableArena) {
let mut y = Vec::new();
for term in &self.terms {
let value = variable_arena
.get(*term)
.map(|variable| variable.get_value())
.unwrap();
y.push(value);
}
let mut fy = Vec::new();
for term in &self.terms {
let value = variable_arena
.get(*term)
.map(|variable| variable.get_message(self.id))
.unwrap();
fy.push(value);
}
let (y, fy) = self.terms
.iter()
.map(|term| {
variable_arena
.get(*term)
.map(|variable| (
variable.get_value(),
variable.get_message(self.id)
))
.unwrap()
})
.unzip();
self.internal_update(variable_arena, self.sum, y, fy, &self.coeffs);
}
pub fn update_term(&self, variable_arena: &mut VariableArena, index: usize) {
let size = self.coeffs.len();
let idx_term = self.terms[index];
let idx_coeff = self.coeffs[index];
let mut a = vec![0.0; size];
let a = self.coeffs
.iter()
.enumerate()
.map(|(i, coeff)| {
if i == index {
1.0 / idx_coeff
} else {
-coeff / idx_coeff
}
})
.collect::<Vec<_>>();
for i in 0..size {
if i != index {
a[i] = -self.coeffs[i] / idx_coeff;
}
}
let (y, fy) = self.terms
.iter()
.enumerate()
.map(|(i, term)| {
let variable = if i == index {
self.sum
} else {
*term
};
a[index] = 1.0 / idx_coeff;
let idx_term = self.terms[index];
let mut y = Vec::new();
let mut fy = Vec::new();
let mut v = self.terms.clone();
v[index] = self.sum;
for term in &v {
let value = variable_arena
.get(*term)
.map(|variable| variable.get_value())
.unwrap();
y.push(value);
let value = variable_arena
.get(*term)
.map(|variable| variable.get_message(self.id))
.unwrap();
fy.push(value);
}
variable_arena
.get(variable)
.map(|variable| (
variable.get_value(),
variable.get_message(self.id)
))
.unwrap()
})
.unzip();
self.internal_update(variable_arena, idx_term, y, fy, &a);
}
@@ -356,14 +328,12 @@ impl TruncateFactor {
}
pub fn update(&self, variable_arena: &mut VariableArena) {
let x = variable_arena
let (x, fx) = variable_arena
.get(self.variable)
.map(|variable| variable.get_value())
.unwrap();
let fx = variable_arena
.get_mut(self.variable)
.map(|variable| variable.get_message(self.id))
.map(|variable| (
variable.get_value(),
variable.get_message(self.id)
))
.unwrap();
let c = x.pi() - fx.pi();