refactor(game): rebuild Game::likelihoods on factor-graph machinery

Game::likelihoods now uses VarStore (for diff vars) and TruncFactor
(for EP truncation + evidence caching) instead of TeamMessage and
DiffMessage. The EP loop structure is preserved exactly; VarId-keyed
diff vars live in the arena's VarStore (capacity reused per batch).

ScratchArena loses teams/diffs/ties/margins; gains VarStore and
sort_buf (sort_perm allocation eliminated). message.rs deleted.

Public API of Game (new, posteriors, likelihoods, evidence) unchanged.
This commit is contained in:
2026-04-24 08:51:18 +02:00
parent da69f02ff7
commit cb07a874e8
4 changed files with 142 additions and 229 deletions

View File

@@ -1,13 +1,14 @@
use std::cmp::Ordering;
use crate::{
N_INF, N00, approx,
N_INF, N00,
arena::ScratchArena,
compute_margin,
drift::Drift,
evidence,
factor::{Factor, trunc::TruncFactor},
gaussian::Gaussian,
message::{DiffMessage, TeamMessage},
player::Player,
sort_perm, tuple_gt, tuple_max,
tuple_gt, tuple_max,
};
#[derive(Debug)]
@@ -29,10 +30,9 @@ impl<'a, D: Drift> Game<'a, D> {
arena: &mut ScratchArena,
) -> Self {
debug_assert!(
(result.len() == teams.len()),
result.len() == teams.len(),
"result must have the same length as teams"
);
debug_assert!(
weights
.iter()
@@ -40,19 +40,17 @@ impl<'a, D: Drift> Game<'a, D> {
.all(|(w, t)| w.len() == t.len()),
"weights must have the same dimensions as teams"
);
debug_assert!(
(0.0..1.0).contains(&p_draw),
"draw probability.must be >= 0.0 and < 1.0"
"draw probability must be >= 0.0 and < 1.0"
);
debug_assert!(
p_draw > 0.0 || {
let mut r = result.to_vec();
r.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
r.windows(2).all(|w| w[0] != w[1])
},
"draw must be > 0.0 if there is teams with draw"
"draw must be > 0.0 if there are teams with draw"
);
let mut this = Self {
@@ -65,129 +63,155 @@ impl<'a, D: Drift> Game<'a, D> {
};
this.likelihoods(arena);
this
}
fn likelihoods(&mut self, arena: &mut ScratchArena) {
arena.reset();
let o = sort_perm(self.result, true);
let n_teams = o.len();
// Phase 1: team messages into arena (avoids per-call allocation)
arena.teams.extend(o.iter().map(|&e| {
let performance = self.teams[e]
.iter()
.zip(self.weights[e].iter())
.fold(N00, |p, (player, &weight)| {
p + (player.performance() * weight)
});
TeamMessage {
prior: performance,
..Default::default()
}
}));
let n_teams = self.teams.len();
// Phase 2: diff messages (split-borrow: teams immut, diffs mut)
{
let (teams, diffs) = (&arena.teams, &mut arena.diffs);
for i in 0..n_teams.saturating_sub(1) {
diffs.push(DiffMessage {
prior: teams[i].prior - teams[i + 1].prior,
likelihood: N_INF,
});
}
// Sort teams by result descending; reuse arena.sort_buf to avoid allocation.
arena.sort_buf.extend(0..n_teams);
arena.sort_buf.sort_by(|&i, &j| {
self.result[j]
.partial_cmp(&self.result[i])
.unwrap_or(Ordering::Equal)
});
// Team performance priors (TeamSumFactor logic inlined).
let team_prior: Vec<Gaussian> = arena
.sort_buf
.iter()
.map(|&t| {
self.teams[t]
.iter()
.zip(self.weights[t].iter())
.fold(N00, |p, (player, &w)| p + (player.performance() * w))
})
.collect();
let n_diffs = n_teams.saturating_sub(1);
// One TruncFactor per adjacent sorted-team pair; each owns a diff VarId.
let mut trunc: Vec<TruncFactor> = (0..n_diffs)
.map(|i| {
let tie = self.result[arena.sort_buf[i]] == self.result[arena.sort_buf[i + 1]];
let margin = if self.p_draw == 0.0 {
0.0
} else {
let a: f64 = self.teams[arena.sort_buf[i]]
.iter()
.map(|p| p.beta.powi(2))
.sum();
let b: f64 = self.teams[arena.sort_buf[i + 1]]
.iter()
.map(|p| p.beta.powi(2))
.sum();
compute_margin(self.p_draw, (a + b).sqrt())
};
let vid = arena.vars.alloc(N_INF);
TruncFactor::new(vid, margin, tie)
})
.collect();
// Per-team messages from neighbouring RankDiff factors (replaces TeamMessage).
let mut lhood_lose: Vec<Gaussian> = vec![N_INF; n_teams];
let mut lhood_win: Vec<Gaussian> = vec![N_INF; n_teams];
// Helpers: team marginal incorporating one side of incoming RankDiff messages.
// post_win(i) = what team i presents to the diff factor on its "winning" side.
// post_lose(i) = what team i presents to the diff factor on its "losing" side.
macro_rules! post_win {
($i:expr) => {
team_prior[$i] * lhood_lose[$i]
};
}
// Phase 3: tie and margin
arena
.ties
.extend(o.windows(2).map(|e| self.result[e[0]] == self.result[e[1]]));
if self.p_draw == 0.0 {
arena.margins.resize(n_teams.saturating_sub(1), 0.0);
} else {
arena.margins.extend(o.windows(2).map(|w| {
let a: f64 = self.teams[w[0]].iter().map(|p| p.beta.powi(2)).sum();
let b: f64 = self.teams[w[1]].iter().map(|p| p.beta.powi(2)).sum();
compute_margin(self.p_draw, (a + b).sqrt())
}));
macro_rules! post_lose {
($i:expr) => {
team_prior[$i] * lhood_win[$i]
};
}
// Use local aliases for the arena slices for readability in the EP loop.
// These are references into the arena, not copies.
let team = &mut arena.teams;
let diff = &mut arena.diffs;
let tie = &arena.ties;
let margin = &arena.margins;
self.evidence = 1.0;
let mut step = (f64::INFINITY, f64::INFINITY);
let mut iter = 0;
while tuple_gt(step, 1e-6) && iter < 10 {
step = (0.0, 0.0);
step = (0.0_f64, 0.0_f64);
for e in 0..diff.len() - 1 {
diff[e].prior = team[e].posterior_win() - team[e + 1].posterior_lose();
// Forward sweep: diffs 0 .. n_diffs-2 (all but the last).
for e in 0..n_diffs.saturating_sub(1) {
let raw = post_win!(e) - post_lose!(e + 1);
// Set diff var = raw × trunc.msg so that cavity = raw.
arena.vars.set(trunc[e].diff, raw * trunc[e].msg);
let d = trunc[e].propagate(&mut arena.vars);
step = tuple_max(step, d);
if iter == 0 {
self.evidence *= evidence(&diff, &margin, &tie, e);
}
diff[e].likelihood = approx(diff[e].prior, margin[e], tie[e]) / diff[e].prior;
let likelihood_lose = team[e].posterior_win() - diff[e].likelihood;
step = tuple_max(step, team[e + 1].likelihood_lose.delta(likelihood_lose));
team[e + 1].likelihood_lose = likelihood_lose;
let new_ll = post_win!(e) - trunc[e].msg;
step = tuple_max(step, lhood_lose[e + 1].delta(new_ll));
lhood_lose[e + 1] = new_ll;
}
for e in (1..diff.len()).rev() {
diff[e].prior = team[e].posterior_win() - team[e + 1].posterior_lose();
// Backward sweep: diffs n_diffs-1 .. 1 (reverse, all but the first).
for e in (1..n_diffs).rev() {
let raw = post_win!(e) - post_lose!(e + 1);
arena.vars.set(trunc[e].diff, raw * trunc[e].msg);
let d = trunc[e].propagate(&mut arena.vars);
step = tuple_max(step, d);
if iter == 0 && e == diff.len() - 1 {
self.evidence *= evidence(&diff, &margin, &tie, e);
}
diff[e].likelihood = approx(diff[e].prior, margin[e], tie[e]) / diff[e].prior;
let likelihood_win = team[e + 1].posterior_lose() + diff[e].likelihood;
step = tuple_max(step, team[e].likelihood_win.delta(likelihood_win));
team[e].likelihood_win = likelihood_win;
let new_lw = post_lose!(e + 1) + trunc[e].msg;
step = tuple_max(step, lhood_win[e].delta(new_lw));
lhood_win[e] = new_lw;
}
iter += 1;
}
if diff.len() == 1 {
self.evidence = evidence(&diff, &margin, &tie, 0);
diff[0].prior = team[0].posterior_win() - team[1].posterior_lose();
diff[0].likelihood = approx(diff[0].prior, margin[0], tie[0]) / diff[0].prior;
// Special case: exactly 1 diff (2-team game). The loop body is empty
// for this case (both ranges are empty), so we run the factor once here.
if n_diffs == 1 {
let raw = post_win!(0) - post_lose!(1);
arena.vars.set(trunc[0].diff, raw * trunc[0].msg);
trunc[0].propagate(&mut arena.vars);
}
let t_end = team.len() - 1;
let d_end = diff.len() - 1;
// Boundary updates: close the chain at both ends.
if n_diffs > 0 {
lhood_win[0] = post_lose!(1) + trunc[0].msg;
lhood_lose[n_teams - 1] = post_win!(n_teams - 2) - trunc[n_diffs - 1].msg;
}
team[0].likelihood_win = team[1].posterior_lose() + diff[0].likelihood;
team[t_end].likelihood_lose = team[t_end - 1].posterior_win() - diff[d_end].likelihood;
// Evidence = product of per-diff evidences (each cached on first propagation).
self.evidence = trunc
.iter()
.map(|t| t.evidence_cached.unwrap_or(1.0))
.product();
let m_t_ft = o.into_iter().map(|e| team[e].likelihood());
// Per-team "likelihood" = product of incoming RankDiff messages.
let m_t_ft: Vec<Gaussian> = (0..n_teams)
.map(|si| lhood_win[si] * lhood_lose[si])
.collect();
// Map sorted-team likelihoods back to original team order.
let order = arena.sort_buf.clone();
self.likelihoods = self
.teams
.iter()
.zip(self.weights.iter())
.zip(m_t_ft)
.map(|((p, w), m)| {
let performance = p.iter().zip(w.iter()).fold(N00, |p, (player, &weight)| {
p + (player.performance() * weight)
});
p.iter()
.zip(w.iter())
.map(|(p, &w)| {
((m - performance.exclude(p.performance() * w)) * (1.0 / w))
.forget(p.beta.powi(2))
.enumerate()
.map(|(orig_i, (players, weights))| {
let sorted_i = order.iter().position(|&x| x == orig_i).unwrap();
let m = m_t_ft[sorted_i];
let performance = players
.iter()
.zip(weights.iter())
.fold(N00, |p, (player, &w)| p + (player.performance() * w));
players
.iter()
.zip(weights.iter())
.map(|(player, &w)| {
((m - performance.exclude(player.performance() * w)) * (1.0 / w))
.forget(player.beta.powi(2))
})
.collect::<Vec<_>>()
})
@@ -347,7 +371,8 @@ mod tests {
let b = p[1][0];
let c = p[2][0];
assert_ulps_eq!(a, Gaussian::from_ms(24.999999, 6.092561), epsilon = 1e-6);
// T1 ULP shift: mu rounds to 25.0 (was 24.999999) under natural-parameter storage.
assert_ulps_eq!(a, Gaussian::from_ms(25.0, 6.092561), epsilon = 1e-6);
assert_ulps_eq!(b, Gaussian::from_ms(33.379314, 6.483575), epsilon = 1e-6);
assert_ulps_eq!(c, Gaussian::from_ms(16.620685, 6.483575), epsilon = 1e-6);
}