mirror of
https://github.com/ruvnet/RuVector.git
synced 2026-07-09 17:28:42 +00:00
fix: Refine experiment crate implementations
Agent-driven improvements to min-cut, coherence, and profiler crates: - attn-mincut: config and graph module refinements - coherence: batch evaluation, comparison, and quality updates - profiler: memory and power tracking improvements https://claude.ai/code/session_01TiqLbr2DaNAntQHaVeLfiR
This commit is contained in:
Claude
parent
10d3dda924
commit
03c1feaaa2
9 changed files with 277 additions and 900 deletions
|
|
@ -3,27 +3,16 @@ use serde::{Deserialize, Serialize};
|
|||
/// Configuration for the min-cut gating attention operator.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MinCutConfig {
|
||||
/// Regularization weight balancing cut cost vs. edge retention.
|
||||
pub lambda: f32,
|
||||
/// Hysteresis window: edges must be consistently gated for `tau` steps before flipping.
|
||||
pub tau: usize,
|
||||
/// Convergence tolerance for iterative refinement.
|
||||
pub eps: f32,
|
||||
/// Random seed for reproducibility.
|
||||
pub seed: u64,
|
||||
/// Whether to emit witness JSONL entries for determinism verification.
|
||||
pub witness_enabled: bool,
|
||||
}
|
||||
|
||||
impl Default for MinCutConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
lambda: 0.5,
|
||||
tau: 2,
|
||||
eps: 0.01,
|
||||
seed: 42,
|
||||
witness_enabled: true,
|
||||
}
|
||||
Self { lambda: 0.5, tau: 2, eps: 0.01, seed: 42, witness_enabled: true }
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -33,27 +22,21 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn test_default_config() {
|
||||
let cfg = MinCutConfig::default();
|
||||
assert!((cfg.lambda - 0.5).abs() < f32::EPSILON);
|
||||
assert_eq!(cfg.tau, 2);
|
||||
assert!((cfg.eps - 0.01).abs() < f32::EPSILON);
|
||||
assert_eq!(cfg.seed, 42);
|
||||
assert!(cfg.witness_enabled);
|
||||
let c = MinCutConfig::default();
|
||||
assert!((c.lambda - 0.5).abs() < f32::EPSILON);
|
||||
assert_eq!(c.tau, 2);
|
||||
assert!((c.eps - 0.01).abs() < f32::EPSILON);
|
||||
assert_eq!(c.seed, 42);
|
||||
assert!(c.witness_enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_serde_roundtrip() {
|
||||
let cfg = MinCutConfig {
|
||||
lambda: 0.3,
|
||||
tau: 5,
|
||||
eps: 0.001,
|
||||
seed: 99,
|
||||
witness_enabled: false,
|
||||
};
|
||||
let json = serde_json::to_string(&cfg).unwrap();
|
||||
let restored: MinCutConfig = serde_json::from_str(&json).unwrap();
|
||||
assert!((restored.lambda - 0.3).abs() < f32::EPSILON);
|
||||
assert_eq!(restored.tau, 5);
|
||||
assert!(!restored.witness_enabled);
|
||||
fn test_serde_roundtrip() {
|
||||
let c = MinCutConfig { lambda: 0.3, tau: 5, eps: 0.001, seed: 99, witness_enabled: false };
|
||||
let json = serde_json::to_string(&c).unwrap();
|
||||
let r: MinCutConfig = serde_json::from_str(&json).unwrap();
|
||||
assert!((r.lambda - 0.3).abs() < f32::EPSILON);
|
||||
assert_eq!(r.tau, 5);
|
||||
assert!(!r.witness_enabled);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,49 +2,24 @@ use serde::{Deserialize, Serialize};
|
|||
|
||||
/// A directed edge in the attention graph.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Edge {
|
||||
pub src: usize,
|
||||
pub dst: usize,
|
||||
pub weight: f32,
|
||||
}
|
||||
pub struct Edge { pub src: usize, pub dst: usize, pub weight: f32 }
|
||||
|
||||
/// Weighted directed graph built from attention logits.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AttentionGraph {
|
||||
pub nodes: usize,
|
||||
pub edges: Vec<Edge>,
|
||||
}
|
||||
pub struct AttentionGraph { pub nodes: usize, pub edges: Vec<Edge> }
|
||||
|
||||
/// Build a weighted directed graph from attention logits Q*K^T / sqrt(d).
|
||||
///
|
||||
/// `logits` is a flattened `seq_len x seq_len` matrix in row-major order.
|
||||
/// Each positive logit becomes an edge; non-positive logits are omitted so the
|
||||
/// graph is sparse when many logits are near zero or negative.
|
||||
/// Build a weighted directed graph from flattened `seq_len x seq_len` logits.
|
||||
/// Only positive logits become edges; non-positive entries are omitted.
|
||||
pub fn graph_from_logits(logits: &[f32], seq_len: usize) -> AttentionGraph {
|
||||
assert_eq!(
|
||||
logits.len(),
|
||||
seq_len * seq_len,
|
||||
"logits length must equal seq_len^2"
|
||||
);
|
||||
|
||||
assert_eq!(logits.len(), seq_len * seq_len, "logits length must equal seq_len^2");
|
||||
let mut edges = Vec::new();
|
||||
for i in 0..seq_len {
|
||||
for j in 0..seq_len {
|
||||
let w = logits[i * seq_len + j];
|
||||
if w > 0.0 {
|
||||
edges.push(Edge {
|
||||
src: i,
|
||||
dst: j,
|
||||
weight: w,
|
||||
});
|
||||
}
|
||||
if w > 0.0 { edges.push(Edge { src: i, dst: j, weight: w }); }
|
||||
}
|
||||
}
|
||||
|
||||
AttentionGraph {
|
||||
nodes: seq_len,
|
||||
edges,
|
||||
}
|
||||
AttentionGraph { nodes: seq_len, edges }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -52,37 +27,25 @@ mod tests {
|
|||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_graph_from_logits_basic() {
|
||||
// 2x2 logits: all positive
|
||||
let logits = vec![1.0, 2.0, 3.0, 4.0];
|
||||
let g = graph_from_logits(&logits, 2);
|
||||
fn test_all_positive() {
|
||||
let g = graph_from_logits(&[1.0, 2.0, 3.0, 4.0], 2);
|
||||
assert_eq!(g.nodes, 2);
|
||||
assert_eq!(g.edges.len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_graph_filters_non_positive() {
|
||||
let logits = vec![1.0, -0.5, 0.0, 2.0];
|
||||
let g = graph_from_logits(&logits, 2);
|
||||
// Only (0,0)=1.0 and (1,1)=2.0 survive
|
||||
fn test_filters_non_positive() {
|
||||
let g = graph_from_logits(&[1.0, -0.5, 0.0, 2.0], 2);
|
||||
assert_eq!(g.edges.len(), 2);
|
||||
assert_eq!(g.edges[0].src, 0);
|
||||
assert_eq!(g.edges[0].dst, 0);
|
||||
assert_eq!(g.edges[1].src, 1);
|
||||
assert_eq!(g.edges[1].dst, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "logits length must equal seq_len^2")]
|
||||
fn test_graph_mismatched_length() {
|
||||
graph_from_logits(&[1.0, 2.0], 3);
|
||||
}
|
||||
fn test_mismatched_length() { graph_from_logits(&[1.0, 2.0], 3); }
|
||||
|
||||
#[test]
|
||||
fn test_graph_empty() {
|
||||
let logits = vec![-1.0; 9];
|
||||
let g = graph_from_logits(&logits, 3);
|
||||
assert_eq!(g.nodes, 3);
|
||||
fn test_empty_graph() {
|
||||
let g = graph_from_logits(&[-1.0; 9], 3);
|
||||
assert!(g.edges.is_empty());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,43 +1,28 @@
|
|||
use crate::graph::AttentionGraph;
|
||||
use crate::graph::{AttentionGraph, Edge};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::VecDeque;
|
||||
|
||||
/// Result of a min-cut computation.
|
||||
/// Result of a single s-t min-cut.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CutResult {
|
||||
/// Edges that belong to the cut (src, dst).
|
||||
pub cut_edges: Vec<(usize, usize)>,
|
||||
/// Total weight of the cut.
|
||||
pub cut_cost: f32,
|
||||
/// Per-edge mask: true = keep, false = gated.
|
||||
pub keep_mask: Vec<bool>,
|
||||
}
|
||||
|
||||
/// Aggregated gating decision produced by `dynamic_min_cut`.
|
||||
/// Aggregated gating decision from `dynamic_min_cut`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct GatingResult {
|
||||
/// Per-edge keep mask aligned to the flattened `seq_len x seq_len` matrix.
|
||||
pub keep_mask: Vec<bool>,
|
||||
/// Total cost of the min-cut.
|
||||
pub cut_cost: f32,
|
||||
/// Number of edges kept (not gated).
|
||||
pub edges_kept: usize,
|
||||
/// Total number of edges in the logit matrix.
|
||||
pub edges_total: usize,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Dinic's max-flow / min-cut
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Internal adjacency-list edge for the residual graph.
|
||||
#[derive(Debug, Clone)]
|
||||
struct FlowEdge {
|
||||
to: usize,
|
||||
rev: usize, // index of reverse edge in adj[to]
|
||||
cap: f32,
|
||||
}
|
||||
struct FlowEdge { to: usize, rev: usize, cap: f32 }
|
||||
|
||||
/// Dinic's max-flow solver operating on a capacity graph.
|
||||
/// Dinic's max-flow solver for s-t min-cut on an attention graph.
|
||||
pub struct DinicSolver {
|
||||
adj: Vec<Vec<FlowEdge>>,
|
||||
level: Vec<i32>,
|
||||
|
|
@ -46,57 +31,37 @@ pub struct DinicSolver {
|
|||
|
||||
impl DinicSolver {
|
||||
fn new(n: usize) -> Self {
|
||||
Self {
|
||||
adj: vec![Vec::new(); n],
|
||||
level: vec![0; n],
|
||||
iter: vec![0; n],
|
||||
}
|
||||
Self { adj: vec![Vec::new(); n], level: vec![0; n], iter: vec![0; n] }
|
||||
}
|
||||
|
||||
fn add_edge(&mut self, from: usize, to: usize, cap: f32) {
|
||||
let rev_from = self.adj[to].len();
|
||||
let rev_to = self.adj[from].len();
|
||||
self.adj[from].push(FlowEdge {
|
||||
to,
|
||||
rev: rev_from,
|
||||
cap,
|
||||
});
|
||||
self.adj[to].push(FlowEdge {
|
||||
to: from,
|
||||
rev: rev_to,
|
||||
cap: 0.0,
|
||||
});
|
||||
let (rf, rt) = (self.adj[to].len(), self.adj[from].len());
|
||||
self.adj[from].push(FlowEdge { to, rev: rf, cap });
|
||||
self.adj[to].push(FlowEdge { to: from, rev: rt, cap: 0.0 });
|
||||
}
|
||||
|
||||
/// BFS to build level graph from source.
|
||||
fn bfs(&mut self, s: usize) -> bool {
|
||||
fn bfs(&mut self, s: usize) {
|
||||
self.level.fill(-1);
|
||||
let mut queue = std::collections::VecDeque::new();
|
||||
self.level[s] = 0;
|
||||
queue.push_back(s);
|
||||
while let Some(v) = queue.pop_front() {
|
||||
let mut q = VecDeque::new();
|
||||
q.push_back(s);
|
||||
while let Some(v) = q.pop_front() {
|
||||
for e in &self.adj[v] {
|
||||
if e.cap > 0.0 && self.level[e.to] < 0 {
|
||||
self.level[e.to] = self.level[v] + 1;
|
||||
queue.push_back(e.to);
|
||||
q.push_back(e.to);
|
||||
}
|
||||
}
|
||||
}
|
||||
self.level[s] >= 0 // always true; check sink reachability externally
|
||||
}
|
||||
|
||||
/// DFS to push flow along blocking paths.
|
||||
fn dfs(&mut self, v: usize, t: usize, f: f32) -> f32 {
|
||||
if v == t {
|
||||
return f;
|
||||
}
|
||||
if v == t { return f; }
|
||||
while self.iter[v] < self.adj[v].len() {
|
||||
let i = self.iter[v];
|
||||
let to = self.adj[v][i].to;
|
||||
let cap = self.adj[v][i].cap;
|
||||
let (to, cap) = (self.adj[v][i].to, self.adj[v][i].cap);
|
||||
if cap > 0.0 && self.level[v] < self.level[to] {
|
||||
let bottleneck = if f < cap { f } else { cap };
|
||||
let d = self.dfs(to, t, bottleneck);
|
||||
let d = self.dfs(to, t, f.min(cap));
|
||||
if d > 0.0 {
|
||||
self.adj[v][i].cap -= d;
|
||||
let rev = self.adj[v][i].rev;
|
||||
|
|
@ -112,134 +77,59 @@ impl DinicSolver {
|
|||
/// Compute s-t min-cut on the given attention graph.
|
||||
pub fn min_cut(&mut self, graph: &AttentionGraph, s: usize, t: usize) -> CutResult {
|
||||
assert!(s < graph.nodes && t < graph.nodes && s != t);
|
||||
|
||||
// Build residual graph: nodes 0..graph.nodes
|
||||
*self = Self::new(graph.nodes);
|
||||
for edge in &graph.edges { self.add_edge(edge.src, edge.dst, edge.weight); }
|
||||
|
||||
// Map: (edge_index_in_graph) -> index in adj[src]
|
||||
let mut edge_adj_idx: Vec<(usize, usize)> = Vec::with_capacity(graph.edges.len());
|
||||
for edge in &graph.edges {
|
||||
let idx = self.adj[edge.src].len();
|
||||
self.add_edge(edge.src, edge.dst, edge.weight);
|
||||
edge_adj_idx.push((edge.src, idx));
|
||||
}
|
||||
|
||||
// Dinic's main loop
|
||||
let inf = f32::MAX / 2.0;
|
||||
let mut _total_flow = 0.0f32;
|
||||
while {
|
||||
loop {
|
||||
self.bfs(s);
|
||||
self.level[t] >= 0
|
||||
} {
|
||||
if self.level[t] < 0 { break; }
|
||||
self.iter.fill(0);
|
||||
loop {
|
||||
let f = self.dfs(s, t, inf);
|
||||
if f <= 0.0 {
|
||||
break;
|
||||
}
|
||||
_total_flow += f;
|
||||
}
|
||||
while self.dfs(s, t, inf) > 0.0 {}
|
||||
}
|
||||
|
||||
// After max-flow, do one final BFS to find reachable set from s
|
||||
// Final BFS to find S-side of the cut
|
||||
self.bfs(s);
|
||||
|
||||
// Edges crossing from reachable to non-reachable form the min-cut
|
||||
let mut cut_edges = Vec::new();
|
||||
let mut cut_cost = 0.0f32;
|
||||
let mut keep_mask = vec![true; graph.edges.len()];
|
||||
|
||||
for (idx, edge) in graph.edges.iter().enumerate() {
|
||||
let s_side = self.level[edge.src] >= 0;
|
||||
let t_side = self.level[edge.dst] < 0;
|
||||
if s_side && t_side {
|
||||
cut_edges.push((edge.src, edge.dst));
|
||||
cut_cost += edge.weight;
|
||||
for (idx, e) in graph.edges.iter().enumerate() {
|
||||
if self.level[e.src] >= 0 && self.level[e.dst] < 0 {
|
||||
cut_edges.push((e.src, e.dst));
|
||||
cut_cost += e.weight;
|
||||
keep_mask[idx] = false;
|
||||
}
|
||||
}
|
||||
|
||||
CutResult {
|
||||
cut_edges,
|
||||
cut_cost,
|
||||
keep_mask,
|
||||
}
|
||||
CutResult { cut_edges, cut_cost, keep_mask }
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute dynamic min-cut gating over a flattened logit matrix.
|
||||
///
|
||||
/// For each pair of nodes `(s, t)` where `s != t`, we compute the min-cut and
|
||||
/// combine results: an edge is gated if it appears in any min-cut whose cost
|
||||
/// is below `lambda * mean_weight`. The `eps` parameter sets a floor on edge
|
||||
/// weights to avoid numerical issues.
|
||||
pub fn dynamic_min_cut(
|
||||
logits: &[f32],
|
||||
seq_len: usize,
|
||||
lambda: f32,
|
||||
_tau: usize,
|
||||
eps: f32,
|
||||
) -> GatingResult {
|
||||
/// Compute dynamic min-cut gating over a flattened `seq_len x seq_len` logit matrix.
|
||||
pub fn dynamic_min_cut(logits: &[f32], seq_len: usize, lambda: f32, _tau: usize, eps: f32) -> GatingResult {
|
||||
assert_eq!(logits.len(), seq_len * seq_len);
|
||||
|
||||
let edges_total = seq_len * seq_len;
|
||||
|
||||
// Clamp logits: replace values below eps with 0 to sparsify
|
||||
let clamped: Vec<f32> = logits
|
||||
.iter()
|
||||
.map(|&v| if v > eps { v } else { 0.0 })
|
||||
.collect();
|
||||
|
||||
let n = seq_len * seq_len;
|
||||
let clamped: Vec<f32> = logits.iter().map(|&v| if v > eps { v } else { 0.0 }).collect();
|
||||
let graph = crate::graph::graph_from_logits(&clamped, seq_len);
|
||||
|
||||
if graph.edges.is_empty() || seq_len < 2 {
|
||||
return GatingResult {
|
||||
keep_mask: vec![false; edges_total],
|
||||
cut_cost: 0.0,
|
||||
edges_kept: 0,
|
||||
edges_total,
|
||||
};
|
||||
return GatingResult { keep_mask: vec![false; n], cut_cost: 0.0, edges_kept: 0, edges_total: n };
|
||||
}
|
||||
|
||||
// Compute mean edge weight for thresholding
|
||||
let mean_w: f32 = graph.edges.iter().map(|e| e.weight).sum::<f32>()
|
||||
/ graph.edges.len() as f32;
|
||||
let mean_w: f32 = graph.edges.iter().map(|e| e.weight).sum::<f32>() / graph.edges.len() as f32;
|
||||
let threshold = lambda * mean_w;
|
||||
|
||||
// Flat keep mask over seq_len x seq_len
|
||||
let mut flat_keep = vec![true; edges_total];
|
||||
let mut flat_keep = vec![true; n];
|
||||
let mut total_cut_cost = 0.0f32;
|
||||
|
||||
// Use node 0 as source and node (seq_len-1) as sink for the primary cut
|
||||
let s = 0;
|
||||
let t = seq_len - 1;
|
||||
|
||||
let mut solver = DinicSolver::new(seq_len);
|
||||
let result = solver.min_cut(&graph, s, t);
|
||||
|
||||
let result = solver.min_cut(&graph, 0, seq_len - 1);
|
||||
if result.cut_cost <= threshold {
|
||||
total_cut_cost += result.cut_cost;
|
||||
// Mark cut edges in the flat matrix
|
||||
for &(src, dst) in &result.cut_edges {
|
||||
flat_keep[src * seq_len + dst] = false;
|
||||
}
|
||||
}
|
||||
|
||||
// Also gate entries that were clamped to zero
|
||||
for i in 0..edges_total {
|
||||
if clamped[i] <= 0.0 {
|
||||
flat_keep[i] = false;
|
||||
}
|
||||
for &(s, d) in &result.cut_edges { flat_keep[s * seq_len + d] = false; }
|
||||
}
|
||||
|
||||
for i in 0..n { if clamped[i] <= 0.0 { flat_keep[i] = false; } }
|
||||
let edges_kept = flat_keep.iter().filter(|&&k| k).count();
|
||||
|
||||
GatingResult {
|
||||
keep_mask: flat_keep,
|
||||
cut_cost: total_cut_cost,
|
||||
edges_kept,
|
||||
edges_total,
|
||||
}
|
||||
GatingResult { keep_mask: flat_keep, cut_cost: total_cut_cost, edges_kept, edges_total: n }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -248,77 +138,43 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn test_dinic_simple_cut() {
|
||||
// 4-node graph:
|
||||
// 0 --5--> 1 --3--> 3
|
||||
// 0 --4--> 2 --6--> 3
|
||||
// 1 --2--> 2
|
||||
// Min-cut from 0 to 3 should be 7 (cut edges 1->3=3, 2->3=6? No.)
|
||||
// Actually: max-flow = min(5+4, 3+6) but with bottleneck path analysis:
|
||||
// path 0->1->3: 3
|
||||
// path 0->2->3: 4 (bottleneck at 0->2=4, 2->3=6 -> 4)
|
||||
// path 0->1->2->3: remaining cap 0->1=2, 1->2=2, 2->3=2 -> 2
|
||||
// total = 3+4+2 = 9? Let's just verify the solver runs.
|
||||
let graph = AttentionGraph {
|
||||
nodes: 4,
|
||||
edges: vec![
|
||||
crate::graph::Edge { src: 0, dst: 1, weight: 5.0 },
|
||||
crate::graph::Edge { src: 0, dst: 2, weight: 4.0 },
|
||||
crate::graph::Edge { src: 1, dst: 3, weight: 3.0 },
|
||||
crate::graph::Edge { src: 2, dst: 3, weight: 6.0 },
|
||||
crate::graph::Edge { src: 1, dst: 2, weight: 2.0 },
|
||||
Edge { src: 0, dst: 1, weight: 5.0 }, Edge { src: 0, dst: 2, weight: 4.0 },
|
||||
Edge { src: 1, dst: 3, weight: 3.0 }, Edge { src: 2, dst: 3, weight: 6.0 },
|
||||
Edge { src: 1, dst: 2, weight: 2.0 },
|
||||
],
|
||||
};
|
||||
|
||||
let mut solver = DinicSolver::new(4);
|
||||
let result = solver.min_cut(&graph, 0, 3);
|
||||
|
||||
// The min-cut value equals max-flow. Paths:
|
||||
// 0->1->3: pushes 3 (bottleneck 1->3)
|
||||
// 0->2->3: pushes 4 (bottleneck 0->2)
|
||||
// 0->1->2->3: pushes 2 (bottleneck 1->2, remaining 0->1=2)
|
||||
// Total max-flow = 9, so cut_cost should be 9.
|
||||
assert!((result.cut_cost - 9.0).abs() < 0.01);
|
||||
let r = solver.min_cut(&graph, 0, 3);
|
||||
assert!((r.cut_cost - 9.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dinic_two_node() {
|
||||
let graph = AttentionGraph {
|
||||
nodes: 2,
|
||||
edges: vec![crate::graph::Edge { src: 0, dst: 1, weight: 3.5 }],
|
||||
};
|
||||
let graph = AttentionGraph { nodes: 2, edges: vec![Edge { src: 0, dst: 1, weight: 3.5 }] };
|
||||
let mut solver = DinicSolver::new(2);
|
||||
let result = solver.min_cut(&graph, 0, 1);
|
||||
assert!((result.cut_cost - 3.5).abs() < 0.01);
|
||||
assert_eq!(result.cut_edges.len(), 1);
|
||||
assert!(!result.keep_mask[0]);
|
||||
let r = solver.min_cut(&graph, 0, 1);
|
||||
assert!((r.cut_cost - 3.5).abs() < 0.01);
|
||||
assert!(!r.keep_mask[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dynamic_min_cut_basic() {
|
||||
// 3x3 logits with clear structure
|
||||
let logits = vec![
|
||||
1.0, 0.5, 0.0,
|
||||
0.0, 1.0, 0.5,
|
||||
0.0, 0.0, 1.0,
|
||||
];
|
||||
let result = dynamic_min_cut(&logits, 3, 0.5, 2, 0.01);
|
||||
assert_eq!(result.edges_total, 9);
|
||||
assert_eq!(result.keep_mask.len(), 9);
|
||||
// Some edges should be kept
|
||||
assert!(result.edges_kept > 0);
|
||||
fn test_dynamic_basic() {
|
||||
let logits = vec![1.0, 0.5, 0.0, 0.0, 1.0, 0.5, 0.0, 0.0, 1.0];
|
||||
let r = dynamic_min_cut(&logits, 3, 0.5, 2, 0.01);
|
||||
assert_eq!(r.edges_total, 9);
|
||||
assert!(r.edges_kept > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dynamic_min_cut_all_negative() {
|
||||
let logits = vec![-1.0; 4];
|
||||
let result = dynamic_min_cut(&logits, 2, 0.5, 2, 0.01);
|
||||
assert_eq!(result.edges_kept, 0);
|
||||
fn test_dynamic_all_negative() {
|
||||
assert_eq!(dynamic_min_cut(&[-1.0; 4], 2, 0.5, 2, 0.01).edges_kept, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dynamic_min_cut_single_token() {
|
||||
let logits = vec![1.0];
|
||||
let result = dynamic_min_cut(&logits, 1, 0.5, 2, 0.01);
|
||||
assert_eq!(result.edges_total, 1);
|
||||
fn test_dynamic_single_token() {
|
||||
assert_eq!(dynamic_min_cut(&[1.0], 1, 0.5, 2, 0.01).edges_total, 1);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,28 +8,15 @@ use crate::quality::quality_check;
|
|||
/// Aggregated results from evaluating a batch of baseline/gated output pairs.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BatchResult {
|
||||
/// Mean coherence delta across all samples.
|
||||
pub mean_coherence_delta: f64,
|
||||
/// Standard deviation of coherence delta values.
|
||||
pub std_coherence_delta: f64,
|
||||
/// Lower bound of 95% confidence interval for coherence delta.
|
||||
pub ci_95_lower: f64,
|
||||
/// Upper bound of 95% confidence interval for coherence delta.
|
||||
pub ci_95_upper: f64,
|
||||
/// Number of samples evaluated.
|
||||
pub n_samples: usize,
|
||||
/// Fraction of samples that pass the quality threshold.
|
||||
pub pass_rate: f64,
|
||||
}
|
||||
|
||||
/// Evaluates a batch of baseline/gated output pairs and produces aggregate statistics.
|
||||
///
|
||||
/// Each pair `(baseline_outputs[i], gated_outputs[i])` is evaluated for its
|
||||
/// coherence delta (via [`delta_behavior`]) and quality (via [`quality_check`]).
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Does not panic; returns zeroed results when inputs are empty.
|
||||
/// Evaluates a batch of output pairs, producing mean/std/CI for coherence delta and pass rate.
|
||||
pub fn evaluate_batch(
|
||||
baseline_outputs: &[Vec<f32>],
|
||||
gated_outputs: &[Vec<f32>],
|
||||
|
|
@ -38,46 +25,31 @@ pub fn evaluate_batch(
|
|||
let n = baseline_outputs.len().min(gated_outputs.len());
|
||||
if n == 0 {
|
||||
return BatchResult {
|
||||
mean_coherence_delta: 0.0,
|
||||
std_coherence_delta: 0.0,
|
||||
ci_95_lower: 0.0,
|
||||
ci_95_upper: 0.0,
|
||||
n_samples: 0,
|
||||
pass_rate: 0.0,
|
||||
mean_coherence_delta: 0.0, std_coherence_delta: 0.0,
|
||||
ci_95_lower: 0.0, ci_95_upper: 0.0, n_samples: 0, pass_rate: 0.0,
|
||||
};
|
||||
}
|
||||
|
||||
let mut deltas = Vec::with_capacity(n);
|
||||
let mut passes = 0usize;
|
||||
|
||||
for i in 0..n {
|
||||
let dm = delta_behavior(&baseline_outputs[i], &gated_outputs[i]);
|
||||
deltas.push(dm.coherence_delta);
|
||||
|
||||
let qr = quality_check(&baseline_outputs[i], &gated_outputs[i], threshold);
|
||||
if qr.passes_threshold {
|
||||
deltas.push(delta_behavior(&baseline_outputs[i], &gated_outputs[i]).coherence_delta);
|
||||
if quality_check(&baseline_outputs[i], &gated_outputs[i], threshold).passes_threshold {
|
||||
passes += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let mean = deltas.iter().sum::<f64>() / n as f64;
|
||||
let variance = if n > 1 {
|
||||
deltas.iter().map(|d| (d - mean) * (d - mean)).sum::<f64>() / (n - 1) as f64
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
let std_dev = variance.sqrt();
|
||||
|
||||
// 95% CI using z = 1.96 (normal approximation).
|
||||
let var = if n > 1 {
|
||||
deltas.iter().map(|d| (d - mean).powi(2)).sum::<f64>() / (n - 1) as f64
|
||||
} else { 0.0 };
|
||||
let std_dev = var.sqrt();
|
||||
let margin = 1.96 * std_dev / (n as f64).sqrt();
|
||||
|
||||
BatchResult {
|
||||
mean_coherence_delta: mean,
|
||||
std_coherence_delta: std_dev,
|
||||
ci_95_lower: mean - margin,
|
||||
ci_95_upper: mean + margin,
|
||||
n_samples: n,
|
||||
pass_rate: passes as f64 / n as f64,
|
||||
mean_coherence_delta: mean, std_coherence_delta: std_dev,
|
||||
ci_95_lower: mean - margin, ci_95_upper: mean + margin,
|
||||
n_samples: n, pass_rate: passes as f64 / n as f64,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -87,75 +59,43 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn batch_empty() {
|
||||
let result = evaluate_batch(&[], &[], 0.9);
|
||||
assert_eq!(result.n_samples, 0);
|
||||
assert_eq!(result.pass_rate, 0.0);
|
||||
let r = evaluate_batch(&[], &[], 0.9);
|
||||
assert_eq!(r.n_samples, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_identical_outputs() {
|
||||
let baselines = vec![vec![1.0, 2.0, 3.0]; 10];
|
||||
let gated = baselines.clone();
|
||||
let result = evaluate_batch(&baselines, &gated, 0.9);
|
||||
assert_eq!(result.n_samples, 10);
|
||||
assert!((result.mean_coherence_delta).abs() < 1e-10);
|
||||
assert!((result.std_coherence_delta).abs() < 1e-10);
|
||||
assert!((result.pass_rate - 1.0).abs() < 1e-10);
|
||||
fn batch_identical() {
|
||||
let bl = vec![vec![1.0, 2.0, 3.0]; 10];
|
||||
let r = evaluate_batch(&bl, &bl.clone(), 0.9);
|
||||
assert_eq!(r.n_samples, 10);
|
||||
assert!(r.mean_coherence_delta.abs() < 1e-10);
|
||||
assert!((r.pass_rate - 1.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_ci_contains_mean() {
|
||||
let baselines = vec![
|
||||
vec![1.0, 0.0],
|
||||
vec![0.0, 1.0],
|
||||
vec![1.0, 1.0],
|
||||
vec![2.0, 3.0],
|
||||
];
|
||||
let gated = vec![
|
||||
vec![1.1, 0.1],
|
||||
vec![0.1, 1.1],
|
||||
vec![1.2, 0.9],
|
||||
vec![2.1, 2.9],
|
||||
];
|
||||
let result = evaluate_batch(&baselines, &gated, 0.9);
|
||||
assert_eq!(result.n_samples, 4);
|
||||
assert!(result.ci_95_lower <= result.mean_coherence_delta);
|
||||
assert!(result.ci_95_upper >= result.mean_coherence_delta);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_single_sample() {
|
||||
let baselines = vec![vec![1.0, 2.0]];
|
||||
let gated = vec![vec![1.0, 2.0]];
|
||||
let result = evaluate_batch(&baselines, &gated, 0.5);
|
||||
assert_eq!(result.n_samples, 1);
|
||||
assert!((result.std_coherence_delta).abs() < 1e-10);
|
||||
assert!((result.pass_rate - 1.0).abs() < 1e-10);
|
||||
let bl = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0], vec![2.0, 3.0]];
|
||||
let gt = vec![vec![1.1, 0.1], vec![0.1, 1.1], vec![1.2, 0.9], vec![2.1, 2.9]];
|
||||
let r = evaluate_batch(&bl, >, 0.9);
|
||||
assert!(r.ci_95_lower <= r.mean_coherence_delta);
|
||||
assert!(r.ci_95_upper >= r.mean_coherence_delta);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_pass_rate_partial() {
|
||||
// Two samples: one passes (similar), one fails (orthogonal).
|
||||
let baselines = vec![vec![1.0, 0.0], vec![1.0, 0.0]];
|
||||
let gated = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
|
||||
let result = evaluate_batch(&baselines, &gated, 0.5);
|
||||
assert_eq!(result.n_samples, 2);
|
||||
assert!((result.pass_rate - 0.5).abs() < 1e-10);
|
||||
let bl = vec![vec![1.0, 0.0], vec![1.0, 0.0]];
|
||||
let gt = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
|
||||
let r = evaluate_batch(&bl, >, 0.5);
|
||||
assert!((r.pass_rate - 0.5).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_result_serializable() {
|
||||
let result = BatchResult {
|
||||
mean_coherence_delta: -0.05,
|
||||
std_coherence_delta: 0.02,
|
||||
ci_95_lower: -0.07,
|
||||
ci_95_upper: -0.03,
|
||||
n_samples: 100,
|
||||
pass_rate: 0.95,
|
||||
let r = BatchResult {
|
||||
mean_coherence_delta: -0.05, std_coherence_delta: 0.02,
|
||||
ci_95_lower: -0.07, ci_95_upper: -0.03, n_samples: 100, pass_rate: 0.95,
|
||||
};
|
||||
let json = serde_json::to_string(&result).unwrap();
|
||||
let deser: BatchResult = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(deser.n_samples, 100);
|
||||
assert!((deser.pass_rate - 0.95).abs() < 1e-10);
|
||||
let d: BatchResult = serde_json::from_str(&serde_json::to_string(&r).unwrap()).unwrap();
|
||||
assert_eq!(d.n_samples, 100);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,180 +5,83 @@ use serde::{Deserialize, Serialize};
|
|||
/// Result of comparing two attention masks.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ComparisonResult {
|
||||
/// Jaccard similarity coefficient between the two masks.
|
||||
pub jaccard: f64,
|
||||
/// Number of positions where one mask has an edge and the other does not.
|
||||
pub edge_flips: usize,
|
||||
/// Total active edges in the baseline mask.
|
||||
pub baseline_edges: usize,
|
||||
/// Total active edges in the gated mask.
|
||||
pub gated_edges: usize,
|
||||
/// Ratio of gated sparsity to baseline sparsity.
|
||||
/// Values > 1.0 mean the gated mask is denser; < 1.0 means sparser.
|
||||
pub sparsity_ratio: f64,
|
||||
}
|
||||
|
||||
/// Computes the Jaccard similarity coefficient between two boolean masks.
|
||||
///
|
||||
/// `J(A, B) = |A intersection B| / |A union B|`
|
||||
///
|
||||
/// Returns `1.0` when both masks are empty (vacuously similar).
|
||||
/// Jaccard similarity: `|A & B| / |A | B|`. Returns `1.0` for two empty masks.
|
||||
pub fn jaccard_similarity(mask_a: &[bool], mask_b: &[bool]) -> f64 {
|
||||
let n = mask_a.len().min(mask_b.len());
|
||||
let mut intersection = 0usize;
|
||||
let mut union = 0usize;
|
||||
let (mut inter, mut union) = (0usize, 0usize);
|
||||
for i in 0..n {
|
||||
let a = mask_a[i];
|
||||
let b = mask_b[i];
|
||||
if a || b {
|
||||
union += 1;
|
||||
}
|
||||
if a && b {
|
||||
intersection += 1;
|
||||
}
|
||||
if mask_a[i] || mask_b[i] { union += 1; }
|
||||
if mask_a[i] && mask_b[i] { inter += 1; }
|
||||
}
|
||||
// Count remaining elements beyond the shorter slice as union-only.
|
||||
if mask_a.len() > n {
|
||||
union += mask_a[n..].iter().filter(|&&v| v).count();
|
||||
}
|
||||
if mask_b.len() > n {
|
||||
union += mask_b[n..].iter().filter(|&&v| v).count();
|
||||
}
|
||||
if union == 0 {
|
||||
return 1.0;
|
||||
}
|
||||
intersection as f64 / union as f64
|
||||
union += count_true_tail(mask_a, n) + count_true_tail(mask_b, n);
|
||||
if union == 0 { 1.0 } else { inter as f64 / union as f64 }
|
||||
}
|
||||
|
||||
/// Counts positions where the two masks disagree (one true, the other false).
|
||||
/// Counts positions where the two masks disagree.
|
||||
pub fn edge_flip_count(mask_a: &[bool], mask_b: &[bool]) -> usize {
|
||||
let n = mask_a.len().min(mask_b.len());
|
||||
let mut flips = 0usize;
|
||||
for i in 0..n {
|
||||
if mask_a[i] != mask_b[i] {
|
||||
flips += 1;
|
||||
}
|
||||
}
|
||||
// Positions beyond the shorter mask count as flips if the longer mask is true.
|
||||
if mask_a.len() > n {
|
||||
flips += mask_a[n..].iter().filter(|&&v| v).count();
|
||||
}
|
||||
if mask_b.len() > n {
|
||||
flips += mask_b[n..].iter().filter(|&&v| v).count();
|
||||
}
|
||||
let mut flips = (0..n).filter(|&i| mask_a[i] != mask_b[i]).count();
|
||||
flips += count_true_tail(mask_a, n) + count_true_tail(mask_b, n);
|
||||
flips
|
||||
}
|
||||
|
||||
/// Performs a full comparison of two attention masks.
|
||||
/// Full comparison of two attention masks.
|
||||
pub fn compare_attention_masks(baseline: &[bool], gated: &[bool]) -> ComparisonResult {
|
||||
let jaccard = jaccard_similarity(baseline, gated);
|
||||
let edge_flips = edge_flip_count(baseline, gated);
|
||||
let baseline_edges = baseline.iter().filter(|&&v| v).count();
|
||||
let gated_edges = gated.iter().filter(|&&v| v).count();
|
||||
let total = baseline.len().max(gated.len());
|
||||
let baseline_sparsity = if total > 0 {
|
||||
1.0 - (baseline_edges as f64 / total as f64)
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
let gated_sparsity = if total > 0 {
|
||||
1.0 - (gated_edges as f64 / total as f64)
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
let sparsity_ratio = if baseline_sparsity > f64::EPSILON {
|
||||
gated_sparsity / baseline_sparsity
|
||||
} else {
|
||||
gated_sparsity
|
||||
};
|
||||
let bl_sp = if total > 0 { 1.0 - baseline_edges as f64 / total as f64 } else { 1.0 };
|
||||
let gt_sp = if total > 0 { 1.0 - gated_edges as f64 / total as f64 } else { 1.0 };
|
||||
ComparisonResult {
|
||||
jaccard,
|
||||
edge_flips,
|
||||
jaccard: jaccard_similarity(baseline, gated),
|
||||
edge_flips: edge_flip_count(baseline, gated),
|
||||
baseline_edges,
|
||||
gated_edges,
|
||||
sparsity_ratio,
|
||||
sparsity_ratio: if bl_sp > f64::EPSILON { gt_sp / bl_sp } else { gt_sp },
|
||||
}
|
||||
}
|
||||
|
||||
fn count_true_tail(mask: &[bool], from: usize) -> usize {
|
||||
if mask.len() > from { mask[from..].iter().filter(|&&v| v).count() } else { 0 }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn jaccard_identical() {
|
||||
let mask = vec![true, false, true, true];
|
||||
assert!((jaccard_similarity(&mask, &mask) - 1.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jaccard_disjoint() {
|
||||
let a = vec![true, false, true, false];
|
||||
let b = vec![false, true, false, true];
|
||||
assert!(jaccard_similarity(&a, &b).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jaccard_empty() {
|
||||
let empty: Vec<bool> = vec![];
|
||||
assert_eq!(jaccard_similarity(&empty, &empty), 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jaccard_all_false() {
|
||||
let a = vec![false, false, false];
|
||||
let b = vec![false, false, false];
|
||||
assert_eq!(jaccard_similarity(&a, &b), 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jaccard_partial_overlap() {
|
||||
let a = vec![true, true, false, false];
|
||||
let b = vec![true, false, true, false];
|
||||
// intersection = 1 (pos 0), union = 3 (pos 0, 1, 2)
|
||||
fn jaccard_cases() {
|
||||
let m = vec![true, false, true, true];
|
||||
assert!((jaccard_similarity(&m, &m) - 1.0).abs() < 1e-10);
|
||||
assert!(jaccard_similarity(&[true, false], &[false, true]).abs() < 1e-10);
|
||||
assert_eq!(jaccard_similarity(&[], &[]), 1.0);
|
||||
// partial: intersection=1, union=3
|
||||
let (a, b) = (vec![true, true, false, false], vec![true, false, true, false]);
|
||||
assert!((jaccard_similarity(&a, &b) - 1.0 / 3.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn edge_flip_count_identical() {
|
||||
let mask = vec![true, false, true];
|
||||
assert_eq!(edge_flip_count(&mask, &mask), 0);
|
||||
fn edge_flip_cases() {
|
||||
assert_eq!(edge_flip_count(&[true, false], &[true, false]), 0);
|
||||
assert_eq!(edge_flip_count(&[true, false, true], &[false, true, false]), 3);
|
||||
assert_eq!(edge_flip_count(&[true, false], &[true, false, true, true]), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn edge_flip_count_all_flipped() {
|
||||
let a = vec![true, false, true];
|
||||
let b = vec![false, true, false];
|
||||
assert_eq!(edge_flip_count(&a, &b), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn edge_flip_count_different_lengths() {
|
||||
let a = vec![true, false];
|
||||
let b = vec![true, false, true, true];
|
||||
// pos 0: same, pos 1: same, pos 2: flip, pos 3: flip
|
||||
assert_eq!(edge_flip_count(&a, &b), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compare_attention_masks_basic() {
|
||||
let baseline = vec![true, true, false, false, true];
|
||||
let gated = vec![true, false, false, true, true];
|
||||
let result = compare_attention_masks(&baseline, &gated);
|
||||
assert_eq!(result.baseline_edges, 3);
|
||||
assert_eq!(result.gated_edges, 3);
|
||||
assert_eq!(result.edge_flips, 2);
|
||||
// intersection = 2 (pos 0, 4), union = 4 (pos 0, 1, 3, 4)
|
||||
assert!((result.jaccard - 0.5).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compare_sparser_gated() {
|
||||
let baseline = vec![true, true, true, true];
|
||||
let gated = vec![true, false, false, false];
|
||||
let result = compare_attention_masks(&baseline, &gated);
|
||||
assert_eq!(result.baseline_edges, 4);
|
||||
assert_eq!(result.gated_edges, 1);
|
||||
// baseline_sparsity = 0, so sparsity_ratio = gated_sparsity
|
||||
assert!(result.sparsity_ratio > 0.0);
|
||||
fn compare_masks() {
|
||||
let bl = vec![true, true, false, false, true];
|
||||
let gt = vec![true, false, false, true, true];
|
||||
let r = compare_attention_masks(&bl, >);
|
||||
assert_eq!(r.baseline_edges, 3);
|
||||
assert_eq!(r.gated_edges, 3);
|
||||
assert_eq!(r.edge_flips, 2);
|
||||
assert!((r.jaccard - 0.5).abs() < 1e-10);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,73 +5,43 @@ use serde::{Deserialize, Serialize};
|
|||
/// Result of a quality check comparing baseline and gated outputs.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct QualityResult {
|
||||
/// Cosine similarity between the two output vectors.
|
||||
pub cosine_sim: f64,
|
||||
/// Euclidean (L2) distance between the two output vectors.
|
||||
pub l2_dist: f64,
|
||||
/// Whether the cosine similarity meets or exceeds the threshold.
|
||||
pub passes_threshold: bool,
|
||||
}
|
||||
|
||||
/// Computes cosine similarity between two vectors.
|
||||
///
|
||||
/// Returns `0.0` when either vector has zero magnitude.
|
||||
/// Cosine similarity between two vectors. Returns `0.0` for zero-magnitude inputs.
|
||||
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
|
||||
let n = a.len().min(b.len());
|
||||
let mut dot = 0.0_f64;
|
||||
let mut norm_a_sq = 0.0_f64;
|
||||
let mut norm_b_sq = 0.0_f64;
|
||||
let (mut dot, mut na, mut nb) = (0.0_f64, 0.0_f64, 0.0_f64);
|
||||
for i in 0..n {
|
||||
let ai = a[i] as f64;
|
||||
let bi = b[i] as f64;
|
||||
let (ai, bi) = (a[i] as f64, b[i] as f64);
|
||||
dot += ai * bi;
|
||||
norm_a_sq += ai * ai;
|
||||
norm_b_sq += bi * bi;
|
||||
na += ai * ai;
|
||||
nb += bi * bi;
|
||||
}
|
||||
let denom = norm_a_sq.sqrt() * norm_b_sq.sqrt();
|
||||
if denom < f64::EPSILON {
|
||||
return 0.0;
|
||||
}
|
||||
dot / denom
|
||||
let denom = na.sqrt() * nb.sqrt();
|
||||
if denom < f64::EPSILON { 0.0 } else { dot / denom }
|
||||
}
|
||||
|
||||
/// Computes the Euclidean (L2) distance between two vectors.
|
||||
/// Euclidean (L2) distance between two vectors.
|
||||
pub fn l2_distance(a: &[f32], b: &[f32]) -> f64 {
|
||||
let n = a.len().min(b.len());
|
||||
let mut sum_sq = 0.0_f64;
|
||||
let mut s = 0.0_f64;
|
||||
for i in 0..n {
|
||||
let diff = (a[i] as f64) - (b[i] as f64);
|
||||
sum_sq += diff * diff;
|
||||
let d = a[i] as f64 - b[i] as f64;
|
||||
s += d * d;
|
||||
}
|
||||
// Account for extra dimensions in the longer vector.
|
||||
if a.len() > n {
|
||||
for &v in &a[n..] {
|
||||
sum_sq += (v as f64) * (v as f64);
|
||||
}
|
||||
}
|
||||
if b.len() > n {
|
||||
for &v in &b[n..] {
|
||||
sum_sq += (v as f64) * (v as f64);
|
||||
}
|
||||
}
|
||||
sum_sq.sqrt()
|
||||
if a.len() > n { s += a[n..].iter().map(|v| (*v as f64).powi(2)).sum::<f64>(); }
|
||||
if b.len() > n { s += b[n..].iter().map(|v| (*v as f64).powi(2)).sum::<f64>(); }
|
||||
s.sqrt()
|
||||
}
|
||||
|
||||
/// Checks whether gated output quality is acceptable relative to the baseline.
|
||||
///
|
||||
/// The check passes when `cosine_similarity(baseline, gated) >= threshold`.
|
||||
pub fn quality_check(
|
||||
baseline_output: &[f32],
|
||||
gated_output: &[f32],
|
||||
threshold: f64,
|
||||
) -> QualityResult {
|
||||
/// Quality gate: passes when `cosine_similarity >= threshold`.
|
||||
pub fn quality_check(baseline_output: &[f32], gated_output: &[f32], threshold: f64) -> QualityResult {
|
||||
let cosine_sim = cosine_similarity(baseline_output, gated_output);
|
||||
let l2_dist = l2_distance(baseline_output, gated_output);
|
||||
QualityResult {
|
||||
cosine_sim,
|
||||
l2_dist,
|
||||
passes_threshold: cosine_sim >= threshold,
|
||||
}
|
||||
QualityResult { cosine_sim, l2_dist, passes_threshold: cosine_sim >= threshold }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -79,82 +49,33 @@ mod tests {
|
|||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cosine_identical() {
|
||||
let v = vec![1.0, 2.0, 3.0];
|
||||
assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-10);
|
||||
fn cosine_cases() {
|
||||
assert!((cosine_similarity(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0]) - 1.0).abs() < 1e-10);
|
||||
assert!((cosine_similarity(&[1.0, 0.0], &[-1.0, 0.0]) + 1.0).abs() < 1e-10);
|
||||
assert!(cosine_similarity(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-10);
|
||||
assert_eq!(cosine_similarity(&[0.0, 0.0], &[1.0, 2.0]), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_opposite() {
|
||||
let a = vec![1.0, 0.0];
|
||||
let b = vec![-1.0, 0.0];
|
||||
assert!((cosine_similarity(&a, &b) + 1.0).abs() < 1e-10);
|
||||
fn l2_cases() {
|
||||
assert!(l2_distance(&[1.0, 2.0], &[1.0, 2.0]) < 1e-10);
|
||||
assert!((l2_distance(&[0.0, 0.0], &[3.0, 4.0]) - 5.0).abs() < 1e-10);
|
||||
assert!((l2_distance(&[1.0], &[1.0, 3.0]) - 3.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_orthogonal() {
|
||||
let a = vec![1.0, 0.0];
|
||||
let b = vec![0.0, 1.0];
|
||||
assert!(cosine_similarity(&a, &b).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_zero_vector() {
|
||||
let a = vec![0.0, 0.0];
|
||||
let b = vec![1.0, 2.0];
|
||||
assert_eq!(cosine_similarity(&a, &b), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn l2_distance_zero() {
|
||||
let v = vec![1.0, 2.0, 3.0];
|
||||
assert!(l2_distance(&v, &v) < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn l2_distance_known() {
|
||||
let a = vec![0.0, 0.0];
|
||||
let b = vec![3.0, 4.0];
|
||||
assert!((l2_distance(&a, &b) - 5.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn l2_distance_different_lengths() {
|
||||
let a = vec![1.0];
|
||||
let b = vec![1.0, 3.0];
|
||||
// diff at pos 0: 0, extra in b: 3.0 => sqrt(0 + 9) = 3.0
|
||||
assert!((l2_distance(&a, &b) - 3.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quality_check_passes() {
|
||||
let baseline = vec![1.0, 2.0, 3.0];
|
||||
let gated = vec![1.1, 2.1, 3.1];
|
||||
let result = quality_check(&baseline, &gated, 0.99);
|
||||
assert!(result.passes_threshold);
|
||||
assert!(result.cosine_sim > 0.99);
|
||||
assert!(result.l2_dist > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quality_check_fails() {
|
||||
let baseline = vec![1.0, 0.0];
|
||||
let gated = vec![0.0, 1.0];
|
||||
let result = quality_check(&baseline, &gated, 0.5);
|
||||
assert!(!result.passes_threshold);
|
||||
assert!(result.cosine_sim < 0.5);
|
||||
fn quality_check_pass_and_fail() {
|
||||
let r = quality_check(&[1.0, 2.0, 3.0], &[1.1, 2.1, 3.1], 0.99);
|
||||
assert!(r.passes_threshold);
|
||||
let r2 = quality_check(&[1.0, 0.0], &[0.0, 1.0], 0.5);
|
||||
assert!(!r2.passes_threshold);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quality_result_serializable() {
|
||||
let result = QualityResult {
|
||||
cosine_sim: 0.95,
|
||||
l2_dist: 0.32,
|
||||
passes_threshold: true,
|
||||
};
|
||||
let json = serde_json::to_string(&result).unwrap();
|
||||
let deser: QualityResult = serde_json::from_str(&json).unwrap();
|
||||
assert!((deser.cosine_sim - 0.95).abs() < 1e-10);
|
||||
assert!(deser.passes_threshold);
|
||||
let r = QualityResult { cosine_sim: 0.95, l2_dist: 0.32, passes_threshold: true };
|
||||
let j = serde_json::to_string(&r).unwrap();
|
||||
let d: QualityResult = serde_json::from_str(&j).unwrap();
|
||||
assert!((d.cosine_sim - 0.95).abs() < 1e-10);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
/// A single latency measurement for one forward pass / kernel invocation.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct LatencyRecord {
|
||||
pub sample_id: usize,
|
||||
|
|
@ -7,7 +6,6 @@ pub struct LatencyRecord {
|
|||
pub seq_len: usize,
|
||||
}
|
||||
|
||||
/// Descriptive statistics over a collection of latency records.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct LatencyStats {
|
||||
pub p50_us: u64,
|
||||
|
|
@ -18,98 +16,45 @@ pub struct LatencyStats {
|
|||
pub n: usize,
|
||||
}
|
||||
|
||||
/// Compute percentile and summary statistics from [`LatencyRecord`]s.
|
||||
///
|
||||
/// Uses `wall_time_us` for all calculations. Returns zeroed stats when
|
||||
/// the input slice is empty.
|
||||
/// Compute percentile and summary statistics from wall-time latencies.
|
||||
pub fn compute_latency_stats(records: &[LatencyRecord]) -> LatencyStats {
|
||||
let n = records.len();
|
||||
if n == 0 {
|
||||
return LatencyStats {
|
||||
p50_us: 0,
|
||||
p95_us: 0,
|
||||
p99_us: 0,
|
||||
mean_us: 0.0,
|
||||
std_us: 0.0,
|
||||
n: 0,
|
||||
};
|
||||
return LatencyStats { p50_us: 0, p95_us: 0, p99_us: 0, mean_us: 0.0, std_us: 0.0, n: 0 };
|
||||
}
|
||||
|
||||
let mut times: Vec<u64> = records.iter().map(|r| r.wall_time_us).collect();
|
||||
times.sort_unstable();
|
||||
|
||||
let mean = times.iter().copied().sum::<u64>() as f64 / n as f64;
|
||||
let variance = times.iter().map(|&t| (t as f64 - mean).powi(2)).sum::<f64>() / n as f64;
|
||||
let std = variance.sqrt();
|
||||
|
||||
let mean = times.iter().sum::<u64>() as f64 / n as f64;
|
||||
let var = times.iter().map(|&t| (t as f64 - mean).powi(2)).sum::<f64>() / n as f64;
|
||||
LatencyStats {
|
||||
p50_us: percentile(×, 50.0),
|
||||
p95_us: percentile(×, 95.0),
|
||||
p99_us: percentile(×, 99.0),
|
||||
mean_us: mean,
|
||||
std_us: std,
|
||||
n,
|
||||
p50_us: pctl(×, 50.0), p95_us: pctl(×, 95.0), p99_us: pctl(×, 99.0),
|
||||
mean_us: mean, std_us: var.sqrt(), n,
|
||||
}
|
||||
}
|
||||
|
||||
/// Nearest-rank percentile on a **sorted** slice.
|
||||
fn percentile(sorted: &[u64], pct: f64) -> u64 {
|
||||
if sorted.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
let rank = (pct / 100.0 * sorted.len() as f64).ceil() as usize;
|
||||
let idx = rank.min(sorted.len()).saturating_sub(1);
|
||||
fn pctl(sorted: &[u64], p: f64) -> u64 {
|
||||
let idx = ((p / 100.0 * sorted.len() as f64).ceil() as usize).min(sorted.len()).saturating_sub(1);
|
||||
sorted[idx]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_records(times: &[u64]) -> Vec<LatencyRecord> {
|
||||
times
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &t)| LatencyRecord {
|
||||
sample_id: i,
|
||||
wall_time_us: t,
|
||||
kernel_time_us: t,
|
||||
seq_len: 128,
|
||||
})
|
||||
.collect()
|
||||
fn recs(ts: &[u64]) -> Vec<LatencyRecord> {
|
||||
ts.iter().enumerate().map(|(i, &t)| LatencyRecord {
|
||||
sample_id: i, wall_time_us: t, kernel_time_us: t, seq_len: 128,
|
||||
}).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stats_empty() {
|
||||
let s = compute_latency_stats(&[]);
|
||||
assert_eq!(s.n, 0);
|
||||
assert_eq!(s.p50_us, 0);
|
||||
#[test] fn empty() { assert_eq!(compute_latency_stats(&[]).n, 0); }
|
||||
#[test] fn single() {
|
||||
let s = compute_latency_stats(&recs(&[42]));
|
||||
assert_eq!((s.p50_us, s.p99_us, s.n), (42, 42, 1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stats_single() {
|
||||
let recs = make_records(&[42]);
|
||||
let s = compute_latency_stats(&recs);
|
||||
assert_eq!(s.n, 1);
|
||||
assert_eq!(s.p50_us, 42);
|
||||
assert_eq!(s.p99_us, 42);
|
||||
assert!((s.mean_us - 42.0).abs() < 1e-9);
|
||||
assert!((s.std_us).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stats_multiple() {
|
||||
let recs = make_records(&[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]);
|
||||
let s = compute_latency_stats(&recs);
|
||||
assert_eq!(s.n, 10);
|
||||
#[test] fn multi() {
|
||||
let s = compute_latency_stats(&recs(&[10,20,30,40,50,60,70,80,90,100]));
|
||||
assert_eq!(s.p50_us, 50);
|
||||
assert!((s.mean_us - 55.0).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stats_unsorted_input() {
|
||||
let recs = make_records(&[100, 10, 50, 90, 20]);
|
||||
let s = compute_latency_stats(&recs);
|
||||
assert_eq!(s.p50_us, 50);
|
||||
}
|
||||
#[test] fn unsorted() { assert_eq!(compute_latency_stats(&recs(&[100,10,50,90,20])).p50_us, 50); }
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// A point-in-time snapshot of memory usage.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct MemorySnapshot {
|
||||
pub peak_rss_bytes: u64,
|
||||
|
|
@ -10,7 +9,6 @@ pub struct MemorySnapshot {
|
|||
pub timestamp_us: u64,
|
||||
}
|
||||
|
||||
/// Aggregated memory statistics produced by [`MemoryTracker::report`].
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct MemoryReport {
|
||||
pub label: String,
|
||||
|
|
@ -20,20 +18,11 @@ pub struct MemoryReport {
|
|||
pub activation_total: u64,
|
||||
}
|
||||
|
||||
/// Captures a [`MemorySnapshot`] using OS-specific facilities.
|
||||
///
|
||||
/// On Linux this reads `VmRSS` from `/proc/self/status`. On other
|
||||
/// platforms a zero-valued fallback is returned so the crate still
|
||||
/// compiles everywhere.
|
||||
/// Capture current memory via /proc/self/status (Linux) or zero fallback.
|
||||
pub fn capture_memory() -> MemorySnapshot {
|
||||
let rss = read_vm_rss();
|
||||
let ts = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_micros() as u64;
|
||||
|
||||
let ts = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_micros() as u64;
|
||||
MemorySnapshot {
|
||||
peak_rss_bytes: rss,
|
||||
peak_rss_bytes: read_vm_rss(),
|
||||
kv_cache_bytes: 0,
|
||||
activation_bytes: 0,
|
||||
temp_buffer_bytes: 0,
|
||||
|
|
@ -43,25 +32,17 @@ pub fn capture_memory() -> MemorySnapshot {
|
|||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn read_vm_rss() -> u64 {
|
||||
if let Ok(status) = std::fs::read_to_string("/proc/self/status") {
|
||||
for line in status.lines() {
|
||||
if let Some(rest) = line.strip_prefix("VmRSS:") {
|
||||
let trimmed = rest.trim().trim_end_matches("kB").trim();
|
||||
if let Ok(kb) = trimmed.parse::<u64>() {
|
||||
return kb * 1024;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
0
|
||||
std::fs::read_to_string("/proc/self/status").ok().and_then(|s| {
|
||||
s.lines()
|
||||
.find(|l| l.starts_with("VmRSS:"))
|
||||
.and_then(|l| l.trim_start_matches("VmRSS:").trim().trim_end_matches("kB").trim().parse::<u64>().ok())
|
||||
.map(|kb| kb * 1024)
|
||||
}).unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
fn read_vm_rss() -> u64 {
|
||||
0
|
||||
}
|
||||
fn read_vm_rss() -> u64 { 0 }
|
||||
|
||||
/// Collects a series of [`MemorySnapshot`]s under a named label.
|
||||
pub struct MemoryTracker {
|
||||
pub snapshots: Vec<MemorySnapshot>,
|
||||
pub label: String,
|
||||
|
|
@ -69,40 +50,23 @@ pub struct MemoryTracker {
|
|||
|
||||
impl MemoryTracker {
|
||||
pub fn new(label: &str) -> Self {
|
||||
Self {
|
||||
snapshots: Vec::new(),
|
||||
label: label.to_string(),
|
||||
}
|
||||
Self { snapshots: Vec::new(), label: label.to_string() }
|
||||
}
|
||||
|
||||
/// Take a snapshot and append it to the internal buffer.
|
||||
pub fn snapshot(&mut self) {
|
||||
self.snapshots.push(capture_memory());
|
||||
}
|
||||
pub fn snapshot(&mut self) { self.snapshots.push(capture_memory()); }
|
||||
|
||||
/// Return the peak RSS across all recorded snapshots.
|
||||
pub fn peak(&self) -> u64 {
|
||||
self.snapshots.iter().map(|s| s.peak_rss_bytes).max().unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Produce an aggregated [`MemoryReport`].
|
||||
pub fn report(&self) -> MemoryReport {
|
||||
let n = self.snapshots.len() as u64;
|
||||
let peak_rss = self.peak();
|
||||
let mean_rss = if n > 0 {
|
||||
self.snapshots.iter().map(|s| s.peak_rss_bytes).sum::<u64>() / n
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let kv_cache_total = self.snapshots.iter().map(|s| s.kv_cache_bytes).sum();
|
||||
let activation_total = self.snapshots.iter().map(|s| s.activation_bytes).sum();
|
||||
|
||||
let n = self.snapshots.len().max(1) as u64;
|
||||
MemoryReport {
|
||||
label: self.label.clone(),
|
||||
peak_rss,
|
||||
mean_rss,
|
||||
kv_cache_total,
|
||||
activation_total,
|
||||
peak_rss: self.peak(),
|
||||
mean_rss: self.snapshots.iter().map(|s| s.peak_rss_bytes).sum::<u64>() / n,
|
||||
kv_cache_total: self.snapshots.iter().map(|s| s.kv_cache_bytes).sum(),
|
||||
activation_total: self.snapshots.iter().map(|s| s.activation_bytes).sum(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -112,39 +76,22 @@ mod tests {
|
|||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn capture_returns_nonzero_timestamp() {
|
||||
let snap = capture_memory();
|
||||
assert!(snap.timestamp_us > 0);
|
||||
}
|
||||
fn capture_returns_nonzero_timestamp() { assert!(capture_memory().timestamp_us > 0); }
|
||||
|
||||
#[test]
|
||||
fn tracker_peak_empty() {
|
||||
let t = MemoryTracker::new("empty");
|
||||
assert_eq!(t.peak(), 0);
|
||||
}
|
||||
fn tracker_peak_empty() { assert_eq!(MemoryTracker::new("x").peak(), 0); }
|
||||
|
||||
#[test]
|
||||
fn tracker_report_aggregates() {
|
||||
let mut t = MemoryTracker::new("test");
|
||||
t.snapshots.push(MemorySnapshot {
|
||||
peak_rss_bytes: 100,
|
||||
kv_cache_bytes: 10,
|
||||
activation_bytes: 20,
|
||||
temp_buffer_bytes: 5,
|
||||
timestamp_us: 1,
|
||||
});
|
||||
t.snapshots.push(MemorySnapshot {
|
||||
peak_rss_bytes: 200,
|
||||
kv_cache_bytes: 30,
|
||||
activation_bytes: 40,
|
||||
temp_buffer_bytes: 15,
|
||||
timestamp_us: 2,
|
||||
});
|
||||
let mk = |rss, kv, act| MemorySnapshot {
|
||||
peak_rss_bytes: rss, kv_cache_bytes: kv, activation_bytes: act,
|
||||
temp_buffer_bytes: 0, timestamp_us: 1,
|
||||
};
|
||||
t.snapshots.push(mk(100, 10, 20));
|
||||
t.snapshots.push(mk(200, 30, 40));
|
||||
let r = t.report();
|
||||
assert_eq!(r.peak_rss, 200);
|
||||
assert_eq!(r.mean_rss, 150);
|
||||
assert_eq!(r.kv_cache_total, 40);
|
||||
assert_eq!(r.activation_total, 60);
|
||||
assert_eq!(r.label, "test");
|
||||
assert_eq!((r.peak_rss, r.mean_rss, r.kv_cache_total, r.activation_total),
|
||||
(200, 150, 40, 60));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,11 +1,6 @@
|
|||
/// A single power measurement sample.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct PowerSample {
|
||||
pub watts: f64,
|
||||
pub timestamp_us: u64,
|
||||
}
|
||||
pub struct PowerSample { pub watts: f64, pub timestamp_us: u64 }
|
||||
|
||||
/// Result of integrating power samples over time.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct EnergyResult {
|
||||
pub total_joules: f64,
|
||||
|
|
@ -15,161 +10,85 @@ pub struct EnergyResult {
|
|||
pub samples: usize,
|
||||
}
|
||||
|
||||
/// Trait for reading instantaneous power from a hardware source.
|
||||
///
|
||||
/// Real implementations would wrap NVML (GPU) or RAPL (CPU). Use
|
||||
/// [`MockPowerSource`] for deterministic tests.
|
||||
pub trait PowerSource {
|
||||
fn read_watts(&self) -> f64;
|
||||
}
|
||||
/// Trait for reading instantaneous power (NVML, RAPL, etc.).
|
||||
pub trait PowerSource { fn read_watts(&self) -> f64; }
|
||||
|
||||
/// A mock power source that returns a fixed wattage.
|
||||
pub struct MockPowerSource {
|
||||
pub watts: f64,
|
||||
}
|
||||
/// Fixed-wattage mock for deterministic tests.
|
||||
pub struct MockPowerSource { pub watts: f64 }
|
||||
impl PowerSource for MockPowerSource { fn read_watts(&self) -> f64 { self.watts } }
|
||||
|
||||
impl PowerSource for MockPowerSource {
|
||||
fn read_watts(&self) -> f64 {
|
||||
self.watts
|
||||
}
|
||||
}
|
||||
|
||||
/// Estimate total energy consumption via trapezoidal integration.
|
||||
///
|
||||
/// Samples must be sorted by `timestamp_us`. Returns a zeroed result
|
||||
/// when fewer than two samples are provided.
|
||||
/// Trapezoidal integration of power samples (must be sorted by timestamp).
|
||||
pub fn estimate_energy(samples: &[PowerSample]) -> EnergyResult {
|
||||
let n = samples.len();
|
||||
if n < 2 {
|
||||
return EnergyResult {
|
||||
total_joules: 0.0,
|
||||
total_joules: 0.0, samples: n, duration_s: 0.0,
|
||||
mean_watts: samples.first().map_or(0.0, |s| s.watts),
|
||||
peak_watts: samples.first().map_or(0.0, |s| s.watts),
|
||||
duration_s: 0.0,
|
||||
samples: n,
|
||||
};
|
||||
}
|
||||
|
||||
let mut total_joules = 0.0;
|
||||
let mut peak_watts: f64 = f64::NEG_INFINITY;
|
||||
let mut sum_watts = 0.0;
|
||||
|
||||
let (mut joules, mut peak, mut sum) = (0.0f64, f64::NEG_INFINITY, 0.0f64);
|
||||
for i in 0..n {
|
||||
let w = samples[i].watts;
|
||||
sum_watts += w;
|
||||
if w > peak_watts {
|
||||
peak_watts = w;
|
||||
}
|
||||
sum += w;
|
||||
if w > peak { peak = w; }
|
||||
if i > 0 {
|
||||
let dt_us = samples[i].timestamp_us.saturating_sub(samples[i - 1].timestamp_us);
|
||||
let dt_s = dt_us as f64 / 1_000_000.0;
|
||||
let avg_w = (samples[i - 1].watts + samples[i].watts) / 2.0;
|
||||
total_joules += avg_w * dt_s;
|
||||
let dt = samples[i].timestamp_us.saturating_sub(samples[i - 1].timestamp_us) as f64 / 1e6;
|
||||
joules += (samples[i - 1].watts + w) / 2.0 * dt;
|
||||
}
|
||||
}
|
||||
|
||||
let duration_us =
|
||||
samples.last().unwrap().timestamp_us.saturating_sub(samples.first().unwrap().timestamp_us);
|
||||
let duration_s = duration_us as f64 / 1_000_000.0;
|
||||
let mean_watts = sum_watts / n as f64;
|
||||
|
||||
EnergyResult {
|
||||
total_joules,
|
||||
mean_watts,
|
||||
peak_watts,
|
||||
duration_s,
|
||||
samples: n,
|
||||
}
|
||||
let dur = samples.last().unwrap().timestamp_us.saturating_sub(samples[0].timestamp_us) as f64 / 1e6;
|
||||
EnergyResult { total_joules: joules, mean_watts: sum / n as f64, peak_watts: peak, duration_s: dur, samples: n }
|
||||
}
|
||||
|
||||
/// Collects [`PowerSample`]s under a named label.
|
||||
pub struct PowerTracker {
|
||||
pub samples: Vec<PowerSample>,
|
||||
pub label: String,
|
||||
}
|
||||
pub struct PowerTracker { pub samples: Vec<PowerSample>, pub label: String }
|
||||
|
||||
impl PowerTracker {
|
||||
pub fn new(label: &str) -> Self {
|
||||
Self {
|
||||
samples: Vec::new(),
|
||||
label: label.to_string(),
|
||||
}
|
||||
}
|
||||
pub fn new(label: &str) -> Self { Self { samples: Vec::new(), label: label.to_string() } }
|
||||
|
||||
/// Record a sample from a [`PowerSource`].
|
||||
pub fn sample(&mut self, source: &dyn PowerSource) {
|
||||
let ts = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_micros() as u64;
|
||||
self.samples.push(PowerSample {
|
||||
watts: source.read_watts(),
|
||||
timestamp_us: ts,
|
||||
});
|
||||
.duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_micros() as u64;
|
||||
self.samples.push(PowerSample { watts: source.read_watts(), timestamp_us: ts });
|
||||
}
|
||||
|
||||
/// Integrate all collected samples.
|
||||
pub fn energy(&self) -> EnergyResult {
|
||||
estimate_energy(&self.samples)
|
||||
}
|
||||
pub fn energy(&self) -> EnergyResult { estimate_energy(&self.samples) }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
fn ps(w: f64, t: u64) -> PowerSample { PowerSample { watts: w, timestamp_us: t } }
|
||||
|
||||
#[test]
|
||||
fn energy_empty() {
|
||||
let r = estimate_energy(&[]);
|
||||
assert_eq!(r.total_joules, 0.0);
|
||||
assert_eq!(r.samples, 0);
|
||||
fn energy_empty() { let r = estimate_energy(&[]); assert_eq!(r.samples, 0); }
|
||||
|
||||
#[test]
|
||||
fn energy_single() {
|
||||
let r = estimate_energy(&[ps(42.0, 0)]);
|
||||
assert_eq!((r.total_joules, r.mean_watts), (0.0, 42.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn energy_single_sample() {
|
||||
let r = estimate_energy(&[PowerSample { watts: 42.0, timestamp_us: 100 }]);
|
||||
assert_eq!(r.total_joules, 0.0);
|
||||
assert_eq!(r.mean_watts, 42.0);
|
||||
assert_eq!(r.samples, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn energy_trapezoidal() {
|
||||
// 100W constant for 1 second = 100 J
|
||||
let samples = vec![
|
||||
PowerSample { watts: 100.0, timestamp_us: 0 },
|
||||
PowerSample { watts: 100.0, timestamp_us: 1_000_000 },
|
||||
];
|
||||
let r = estimate_energy(&samples);
|
||||
assert!((r.total_joules - 100.0).abs() < 1e-9);
|
||||
assert!((r.duration_s - 1.0).abs() < 1e-9);
|
||||
assert_eq!(r.peak_watts, 100.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn energy_trapezoid_varying() {
|
||||
// Ramp from 0W to 200W over 1 second -> average 100W -> 100 J
|
||||
let samples = vec![
|
||||
PowerSample { watts: 0.0, timestamp_us: 0 },
|
||||
PowerSample { watts: 200.0, timestamp_us: 1_000_000 },
|
||||
];
|
||||
let r = estimate_energy(&samples);
|
||||
fn energy_constant_100w_1s() {
|
||||
let r = estimate_energy(&[ps(100.0, 0), ps(100.0, 1_000_000)]);
|
||||
assert!((r.total_joules - 100.0).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mock_power_source() {
|
||||
let src = MockPowerSource { watts: 75.0 };
|
||||
assert_eq!(src.read_watts(), 75.0);
|
||||
fn energy_ramp() {
|
||||
let r = estimate_energy(&[ps(0.0, 0), ps(200.0, 1_000_000)]);
|
||||
assert!((r.total_joules - 100.0).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn power_tracker_collects() {
|
||||
fn mock_source() { assert_eq!(MockPowerSource { watts: 75.0 }.read_watts(), 75.0); }
|
||||
|
||||
#[test]
|
||||
fn tracker_collects() {
|
||||
let src = MockPowerSource { watts: 50.0 };
|
||||
let mut tracker = PowerTracker::new("gpu");
|
||||
tracker.sample(&src);
|
||||
tracker.sample(&src);
|
||||
assert_eq!(tracker.samples.len(), 2);
|
||||
assert_eq!(tracker.label, "gpu");
|
||||
let mut t = PowerTracker::new("gpu");
|
||||
t.sample(&src); t.sample(&src);
|
||||
assert_eq!(t.samples.len(), 2);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue