diff --git a/src/factor/mod.rs b/src/factor/mod.rs index 5bc76f0..6caf161 100644 --- a/src/factor/mod.rs +++ b/src/factor/mod.rs @@ -7,7 +7,7 @@ use crate::gaussian::Gaussian; /// Variables hold the current Gaussian marginal and are owned by exactly one /// `VarStore`. `VarId` is meaningful only within its owning store. #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] -pub(crate) struct VarId(pub(crate) u32); +pub struct VarId(pub u32); /// Flat storage of variable marginals. /// @@ -15,36 +15,38 @@ pub(crate) struct VarId(pub(crate) u32); /// reused across `Game::new` calls (it lives in the `ScratchArena`); call /// `clear()` before reuse. #[derive(Debug, Default)] -pub(crate) struct VarStore { +pub struct VarStore { pub(crate) marginals: Vec, } impl VarStore { - #[allow(dead_code)] - pub(crate) fn new() -> Self { + pub fn new() -> Self { Self::default() } - pub(crate) fn clear(&mut self) { + pub fn clear(&mut self) { self.marginals.clear(); } - #[allow(dead_code)] - pub(crate) fn len(&self) -> usize { + pub fn len(&self) -> usize { self.marginals.len() } - pub(crate) fn alloc(&mut self, init: Gaussian) -> VarId { + pub fn is_empty(&self) -> bool { + self.marginals.is_empty() + } + + pub fn alloc(&mut self, init: Gaussian) -> VarId { let id = VarId(self.marginals.len() as u32); self.marginals.push(init); id } - pub(crate) fn get(&self, id: VarId) -> Gaussian { + pub fn get(&self, id: VarId) -> Gaussian { self.marginals[id.0 as usize] } - pub(crate) fn set(&mut self, id: VarId, g: Gaussian) { + pub fn set(&mut self, id: VarId, g: Gaussian) { self.marginals[id.0 as usize] = g; } } @@ -54,7 +56,7 @@ impl VarStore { /// Factors hold their own outgoing messages and propagate them by reading /// connected variable marginals from a `VarStore` and writing back updated /// marginals. -pub(crate) trait Factor { +pub trait Factor { /// Update outgoing messages and write back to the var store. /// /// Returns the max delta `(|Δmu|, |Δsigma|)` across writes this @@ -62,7 +64,6 @@ pub(crate) trait Factor { fn propagate(&mut self, vars: &mut VarStore) -> (f64, f64); /// Optional log-evidence contribution. Default 0.0 (no contribution). - #[allow(dead_code)] fn log_evidence(&self, _vars: &VarStore) -> f64 { 0.0 } @@ -73,8 +74,7 @@ pub(crate) trait Factor { /// Using an enum instead of `Box` keeps factor data inline and /// avoids virtual-call overhead in the hot inference loop. #[derive(Debug)] -#[allow(dead_code)] -pub(crate) enum BuiltinFactor { +pub enum BuiltinFactor { TeamSum(team_sum::TeamSumFactor), RankDiff(rank_diff::RankDiffFactor), Trunc(trunc::TruncFactor), @@ -97,9 +97,9 @@ impl Factor for BuiltinFactor { } } -pub(crate) mod rank_diff; -pub(crate) mod team_sum; -pub(crate) mod trunc; +pub mod rank_diff; +pub mod team_sum; +pub mod trunc; #[cfg(test)] mod tests { diff --git a/src/factor/rank_diff.rs b/src/factor/rank_diff.rs index c48bab3..ce64a95 100644 --- a/src/factor/rank_diff.rs +++ b/src/factor/rank_diff.rs @@ -13,11 +13,10 @@ use crate::factor::{Factor, VarId, VarStore}; /// effectively replaced on each propagation. The TruncFactor on the same diff /// var holds the EP-divide message that produces the cavity. #[derive(Debug)] -#[allow(dead_code)] -pub(crate) struct RankDiffFactor { - pub(crate) team_a: VarId, - pub(crate) team_b: VarId, - pub(crate) diff: VarId, +pub struct RankDiffFactor { + pub team_a: VarId, + pub team_b: VarId, + pub diff: VarId, } impl Factor for RankDiffFactor { diff --git a/src/factor/team_sum.rs b/src/factor/team_sum.rs index 33d93dc..a110141 100644 --- a/src/factor/team_sum.rs +++ b/src/factor/team_sum.rs @@ -10,10 +10,9 @@ use crate::{ /// already with beta² noise added via `Rating::performance()`). The factor /// runs once per game and writes the weighted sum to the output var. #[derive(Debug)] -#[allow(dead_code)] -pub(crate) struct TeamSumFactor { - pub(crate) inputs: Vec<(Gaussian, f64)>, - pub(crate) out: VarId, +pub struct TeamSumFactor { + pub inputs: Vec<(Gaussian, f64)>, + pub out: VarId, } impl Factor for TeamSumFactor { diff --git a/src/factor/trunc.rs b/src/factor/trunc.rs index abe3dbb..6090a39 100644 --- a/src/factor/trunc.rs +++ b/src/factor/trunc.rs @@ -11,10 +11,10 @@ use crate::{ /// Stores its outgoing message to the diff variable so the cavity computation /// produces the correct EP message on each propagation. #[derive(Debug)] -pub(crate) struct TruncFactor { - pub(crate) diff: VarId, - pub(crate) margin: f64, - pub(crate) tie: bool, +pub struct TruncFactor { + pub diff: VarId, + pub margin: f64, + pub tie: bool, /// Outgoing message to the diff variable (initial: N_INF, the EP identity). pub(crate) msg: Gaussian, /// Cached evidence (linear, not log) computed from the cavity on first propagation. @@ -22,7 +22,7 @@ pub(crate) struct TruncFactor { } impl TruncFactor { - pub(crate) fn new(diff: VarId, margin: f64, tie: bool) -> Self { + pub fn new(diff: VarId, margin: f64, tie: bool) -> Self { Self { diff, margin, diff --git a/src/factors.rs b/src/factors.rs new file mode 100644 index 0000000..162ca68 --- /dev/null +++ b/src/factors.rs @@ -0,0 +1,13 @@ +//! Factor-graph public API. +//! +//! Power users can construct custom factor graphs via `Game::custom` (T2 +//! minimal; full ergonomics in T4) and drive them with custom `Schedule` +//! implementations. + +pub use crate::{ + factor::{ + BuiltinFactor, Factor, VarId, VarStore, rank_diff::RankDiffFactor, team_sum::TeamSumFactor, + trunc::TruncFactor, + }, + schedule::{EpsilonOrMax, Schedule, ScheduleReport}, +}; diff --git a/src/lib.rs b/src/lib.rs index 63ab8a0..09b7657 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ mod error; mod event; mod event_builder; pub(crate) mod factor; +pub mod factors; mod game; pub mod gaussian; mod history; diff --git a/src/schedule.rs b/src/schedule.rs index 7d1336f..2c98fc1 100644 --- a/src/schedule.rs +++ b/src/schedule.rs @@ -16,8 +16,7 @@ pub struct ScheduleReport { } /// Drives factor propagation to convergence. -#[allow(dead_code)] -pub(crate) trait Schedule { +pub trait Schedule { fn run(&self, factors: &mut [BuiltinFactor], vars: &mut VarStore) -> ScheduleReport; } @@ -26,8 +25,7 @@ pub(crate) trait Schedule { /// Matches the existing `Game::likelihoods` loop bit-for-bit when given the /// same factor layout (TeamSums first, then alternating RankDiff/Trunc pairs). #[derive(Debug, Clone, Copy)] -#[allow(dead_code)] -pub(crate) struct EpsilonOrMax { +pub struct EpsilonOrMax { pub eps: f64, pub max: usize, }