T0 + T1 + T2: engine redesign through new API surface #1

Merged
logaritmisk merged 45 commits from t2-new-api-surface into main 2026-04-24 11:20:04 +00:00
6 changed files with 234 additions and 61 deletions
Showing only changes of commit b1e0fcb817 - Show all commits

44
src/arena.rs Normal file
View File

@@ -0,0 +1,44 @@
use crate::message::{DiffMessage, TeamMessage};
/// Reusable scratch buffers for `Game::likelihoods`.
///
/// The four Vecs previously allocated fresh on every `Game::new` call —
/// `teams`, `diffs`, `ties`, `margins` — are now borrowed from this arena,
/// reset between uses. A `Batch` owns one arena; all events in the slice
/// share it across the convergence iterations.
#[derive(Debug, Default)]
pub struct ScratchArena {
pub(crate) teams: Vec<TeamMessage>,
pub(crate) diffs: Vec<DiffMessage>,
pub(crate) ties: Vec<bool>,
pub(crate) margins: Vec<f64>,
}
impl ScratchArena {
pub fn new() -> Self {
Self::default()
}
#[inline]
pub(crate) fn reset(&mut self) {
self.teams.clear();
self.diffs.clear();
self.ties.clear();
self.margins.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reset_keeps_capacity() {
let mut arena = ScratchArena::new();
arena.teams.push(TeamMessage::default());
let cap = arena.teams.capacity();
arena.reset();
assert_eq!(arena.teams.len(), 0);
assert_eq!(arena.teams.capacity(), cap);
}
}

View File

@@ -3,6 +3,7 @@ use std::collections::HashMap;
use crate::{
Index, N_INF,
agent::Agent,
arena::ScratchArena,
drift::Drift,
game::Game,
gaussian::Gaussian,
@@ -111,6 +112,7 @@ pub struct Batch {
pub(crate) skills: SkillStore,
pub(crate) time: i64,
p_draw: f64,
arena: ScratchArena,
}
impl Batch {
@@ -120,6 +122,7 @@ impl Batch {
skills: SkillStore::new(),
time,
p_draw,
arena: ScratchArena::new(),
}
}
@@ -219,7 +222,7 @@ impl Batch {
let teams = event.within_priors(false, false, &self.skills, agents);
let result = event.outputs();
let g = Game::new(teams, &result, &event.weights, self.p_draw);
let g = Game::new(teams, &result, &event.weights, self.p_draw, &mut self.arena);
for (t, team) in event.teams.iter_mut().enumerate() {
for (i, item) in team.items.iter_mut().enumerate() {
@@ -295,6 +298,9 @@ impl Batch {
forward: bool,
agents: &AgentStore<D>,
) -> f64 {
// log_evidence is infrequent; a local arena avoids needing &mut self.
let mut arena = ScratchArena::new();
if targets.is_empty() {
if online || forward {
self.events
@@ -306,6 +312,7 @@ impl Batch {
&event.outputs(),
&event.weights,
self.p_draw,
&mut arena,
)
.evidence
.ln()
@@ -331,6 +338,7 @@ impl Batch {
&event.outputs(),
&event.weights,
self.p_draw,
&mut arena,
)
.evidence
.ln()

View File

@@ -1,5 +1,7 @@
use crate::{
N_INF, N00, approx, compute_margin,
N_INF, N00, approx,
arena::ScratchArena,
compute_margin,
drift::Drift,
evidence,
gaussian::Gaussian,
@@ -24,6 +26,7 @@ impl<'a, D: Drift> Game<'a, D> {
result: &'a [f64],
weights: &'a [Vec<f64>],
p_draw: f64,
arena: &mut ScratchArena,
) -> Self {
debug_assert!(
(result.len() == teams.len()),
@@ -61,56 +64,62 @@ impl<'a, D: Drift> Game<'a, D> {
evidence: 0.0,
};
this.likelihoods();
this.likelihoods(arena);
this
}
fn likelihoods(&mut self) {
fn likelihoods(&mut self, arena: &mut ScratchArena) {
arena.reset();
let o = sort_perm(self.result, true);
let n_teams = o.len();
let mut team = o
.iter()
.map(|&e| {
let performance = self.teams[e]
.iter()
.zip(self.weights[e].iter())
.fold(N00, |p, (player, &weight)| {
p + (player.performance() * weight)
});
// 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()
}
}));
TeamMessage {
prior: performance,
..Default::default()
}
})
.collect::<Vec<_>>();
// 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,
});
}
}
let mut diff = team
.windows(2)
.map(|w| DiffMessage {
prior: w[0].prior - w[1].prior,
likelihood: N_INF,
})
.collect::<Vec<_>>();
// Phase 3: tie and margin
arena
.ties
.extend(o.windows(2).map(|e| self.result[e[0]] == self.result[e[1]]));
let tie = o
.windows(2)
.map(|e| self.result[e[0]] == self.result[e[1]])
.collect::<Vec<_>>();
let margin = if self.p_draw == 0.0 {
vec![0.0; o.len() - 1]
if self.p_draw == 0.0 {
arena.margins.resize(n_teams.saturating_sub(1), 0.0);
} else {
o.windows(2)
.map(|w| {
let a: f64 = self.teams[w[0]].iter().map(|a| a.beta.powi(2)).sum();
let b: f64 = self.teams[w[1]].iter().map(|a| a.beta.powi(2)).sum();
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())
}));
}
compute_margin(self.p_draw, (a + b).sqrt())
})
.collect::<Vec<_>>()
};
// 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;
@@ -204,7 +213,7 @@ mod tests {
use ::approx::assert_ulps_eq;
use super::*;
use crate::{ConstantDrift, GAMMA, Gaussian, N_INF, Player};
use crate::{ConstantDrift, GAMMA, Gaussian, N_INF, Player, arena::ScratchArena};
#[test]
fn test_1vs1() {
@@ -220,7 +229,13 @@ mod tests {
);
let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 1.0], &w, 0.0);
let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 1.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
let a = p[0][0];
@@ -241,7 +256,13 @@ mod tests {
);
let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 1.0], &w, 0.0);
let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 1.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
let a = p[0][0];
@@ -254,7 +275,13 @@ mod tests {
let t_b = Player::new(Gaussian::from_ms(15.568, 0.51), 1.0, ConstantDrift(0.2125));
let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 1.0], &w, 0.0);
let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 1.0],
&w,
0.0,
&mut ScratchArena::new(),
);
assert_eq!(g.likelihoods[0][0], N_INF);
assert_eq!(g.likelihoods[1][0], N_INF);
@@ -281,7 +308,13 @@ mod tests {
];
let w = [vec![1.0], vec![1.0], vec![1.0]];
let g = Game::new(teams.clone(), &[1.0, 2.0, 0.0], &w, 0.0);
let g = Game::new(
teams.clone(),
&[1.0, 2.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
let a = p[0][0];
@@ -291,7 +324,13 @@ mod tests {
assert_ulps_eq!(b, Gaussian::from_ms(31.311358, 6.698818), epsilon = 1e-6);
let w = [vec![1.0], vec![1.0], vec![1.0]];
let g = Game::new(teams.clone(), &[2.0, 1.0, 0.0], &w, 0.0);
let g = Game::new(
teams.clone(),
&[2.0, 1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
let a = p[0][0];
@@ -301,7 +340,7 @@ mod tests {
assert_ulps_eq!(b, Gaussian::from_ms(25.000000, 6.238469), epsilon = 1e-6);
let w = [vec![1.0], vec![1.0], vec![1.0]];
let g = Game::new(teams, &[1.0, 2.0, 0.0], &w, 0.5);
let g = Game::new(teams, &[1.0, 2.0, 0.0], &w, 0.5, &mut ScratchArena::new());
let p = g.posteriors();
let a = p[0][0];
@@ -327,7 +366,13 @@ mod tests {
);
let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 0.0], &w, 0.25);
let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors();
let a = p[0][0];
@@ -348,7 +393,13 @@ mod tests {
);
let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 0.0], &w, 0.25);
let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors();
let a = p[0][0];
@@ -382,6 +433,7 @@ mod tests {
&[0.0, 0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors();
@@ -417,6 +469,7 @@ mod tests {
&[0.0, 0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors();
@@ -462,7 +515,13 @@ mod tests {
];
let w = [vec![1.0, 1.0], vec![1.0], vec![1.0, 1.0]];
let g = Game::new(vec![t_a, t_b, t_c], &[1.0, 0.0, 0.0], &w, 0.25);
let g = Game::new(
vec![t_a, t_b, t_c],
&[1.0, 0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(p[0][0], Gaussian::from_ms(13.051, 2.864), epsilon = 1e-3);
@@ -489,7 +548,13 @@ mod tests {
)];
let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(
@@ -507,7 +572,13 @@ mod tests {
let w_b = vec![0.7];
let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(
@@ -525,7 +596,13 @@ mod tests {
let w_b = vec![0.7];
let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a, t_b],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(
@@ -554,7 +631,13 @@ mod tests {
)];
let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a, t_b],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(
@@ -583,7 +666,13 @@ mod tests {
)];
let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a, t_b],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(p[0][0], p[1][0], epsilon = 1e-6);
@@ -620,7 +709,13 @@ mod tests {
let w_b = vec![0.9, 0.6];
let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(
@@ -648,7 +743,13 @@ mod tests {
let w_b = vec![0.7, 0.4];
let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(
@@ -676,7 +777,13 @@ mod tests {
let w_b = vec![0.7, 2.4];
let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(
@@ -713,6 +820,7 @@ mod tests {
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let post_2vs1 = g.posteriors();
@@ -720,7 +828,13 @@ mod tests {
let w_b = vec![1.0, 0.0];
let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b.clone()], &[1.0, 0.0], &w, 0.0);
let g = Game::new(
vec![t_a, t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors();
assert_ulps_eq!(p[0][0], post_2vs1[0][0], epsilon = 1e-6);

View File

@@ -436,7 +436,10 @@ mod tests {
use approx::assert_ulps_eq;
use super::*;
use crate::{ConstantDrift, EPSILON, Game, Gaussian, ITERATIONS, IndexMap, P_DRAW, Player};
use crate::{
ConstantDrift, EPSILON, Game, Gaussian, ITERATIONS, IndexMap, P_DRAW, Player,
arena::ScratchArena,
};
#[test]
fn test_init() {
@@ -500,6 +503,7 @@ mod tests {
&[0.0, 1.0],
&w,
P_DRAW,
&mut ScratchArena::new(),
)
.posteriors();
let expected = p[0][0];

View File

@@ -9,6 +9,7 @@ use std::{
pub mod agent;
#[cfg(feature = "approx")]
mod approx;
pub(crate) mod arena;
pub mod batch;
pub mod drift;
mod error;

View File

@@ -1,5 +1,6 @@
use crate::{N_INF, gaussian::Gaussian};
#[derive(Debug)]
pub(crate) struct TeamMessage {
pub(crate) prior: Gaussian,
pub(crate) likelihood_lose: Gaussian,
@@ -67,6 +68,7 @@ impl DrawMessage {
}
}
*/
#[derive(Debug)]
pub(crate) struct DiffMessage {
pub(crate) prior: Gaussian,
pub(crate) likelihood: Gaussian,