From 9836b7b709085d7396f11130517ca301993ea384 Mon Sep 17 00:00:00 2001 From: Anders Olsson Date: Fri, 24 Apr 2026 13:42:05 +0200 Subject: [PATCH] feat(time-slice): compute and maintain color groups; reorder events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TimeSlice gains a color_groups field of type ColorGroups, recomputed whenever events change. After recompute, self.events is physically reordered so color-0 events are first, then color-1, etc. Each color is therefore a contiguous range of indices in self.events — the invariant that Task 6's parallel par_iter_mut exploits. Greedy coloring via crate::color_group::color_greedy; agent indices come from Event::iter_agents. ColorGroups gains a color_range helper that returns the contiguous Range for a given color. Numerical behavior unchanged: async-EP is order-independent at convergence, so event reordering does not affect goldens. Part of T3. Co-Authored-By: Claude Sonnet 4.6 --- src/color_group.rs | 13 ++++++ src/time_slice.rs | 107 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/src/color_group.rs b/src/color_group.rs index af49c20..6add43c 100644 --- a/src/color_group.rs +++ b/src/color_group.rs @@ -46,6 +46,19 @@ impl ColorGroups { pub(crate) fn total_events(&self) -> usize { self.groups.iter().map(|g| g.len()).sum() } + + /// Contiguous index range for one color after events have been reordered + /// into color-contiguous positions by `TimeSlice::recompute_color_groups`. + #[allow(dead_code)] + pub(crate) fn color_range(&self, color_idx: usize) -> std::ops::Range { + let group = &self.groups[color_idx]; + if group.is_empty() { + return 0..0; + } + let start = *group.first().unwrap(); + let end = *group.last().unwrap() + 1; + start..end + } } /// Compute color groups greedily. diff --git a/src/time_slice.rs b/src/time_slice.rs index c1d48fb..bf6008d 100644 --- a/src/time_slice.rs +++ b/src/time_slice.rs @@ -7,6 +7,7 @@ use std::collections::HashMap; use crate::{ Index, N_INF, arena::ScratchArena, + color_group::ColorGroups, drift::Drift, game::Game, gaussian::Gaussian, @@ -84,6 +85,12 @@ pub(crate) struct Event { } impl Event { + pub(crate) fn iter_agents(&self) -> impl Iterator + '_ { + self.teams + .iter() + .flat_map(|t| t.items.iter().map(|it| it.agent)) + } + fn outputs(&self) -> Vec { self.teams .iter() @@ -117,6 +124,7 @@ pub struct TimeSlice { pub(crate) time: T, p_draw: f64, arena: ScratchArena, + pub(crate) color_groups: ColorGroups, } impl TimeSlice { @@ -127,9 +135,44 @@ impl TimeSlice { time, p_draw, arena: ScratchArena::new(), + color_groups: ColorGroups::new(), } } + /// Recompute the color-group partition and reorder `self.events` into + /// color-contiguous ranges. After this call, `self.color_groups.groups[c]` + /// contains a contiguous ascending range of indices in `self.events`. + pub(crate) fn recompute_color_groups(&mut self) { + use crate::color_group::color_greedy; + + let n = self.events.len(); + if n == 0 { + self.color_groups = ColorGroups::new(); + return; + } + + let cg = color_greedy(n, |ev_idx| { + self.events[ev_idx].iter_agents().collect::>() + }); + + let mut reordered: Vec = Vec::with_capacity(n); + let mut new_groups: Vec> = Vec::with_capacity(cg.groups.len()); + let mut taken: Vec> = self.events.drain(..).map(Some).collect(); + + for group in &cg.groups { + let mut new_indices: Vec = Vec::with_capacity(group.len()); + for &old_idx in group { + let ev = taken[old_idx].take().expect("event already taken"); + new_indices.push(reordered.len()); + reordered.push(ev); + } + new_groups.push(new_indices); + } + + self.events = reordered; + self.color_groups = ColorGroups { groups: new_groups }; + } + pub fn add_events>( &mut self, composition: Vec>>, @@ -212,6 +255,7 @@ impl TimeSlice { self.events.extend(events); self.iteration(from, agents); + self.recompute_color_groups(); } pub(crate) fn posteriors(&self) -> HashMap { @@ -662,4 +706,67 @@ mod tests { epsilon = 1e-6 ); } + + #[test] + fn time_slice_color_groups_reorders_events() { + // ev0: [a, b]; ev1: [c, d]; ev2: [a, c] + // Greedy coloring: ev0→c0, ev1→c0 (disjoint), ev2→c1 (overlaps both). + // After recompute_color_groups, physical order is [ev0, ev1, ev2] + // and groups == [[0, 1], [2]]. + let mut index_map = KeyTable::new(); + + let a = index_map.get_or_create("a"); + let b = index_map.get_or_create("b"); + let c = index_map.get_or_create("c"); + let d = index_map.get_or_create("d"); + + let mut agents: CompetitorStore = CompetitorStore::new(); + + for agent in [a, b, c, d] { + agents.insert( + agent, + Competitor { + rating: Rating::new( + Gaussian::from_ms(25.0, 25.0 / 3.0), + 25.0 / 6.0, + ConstantDrift(25.0 / 300.0), + ), + ..Default::default() + }, + ); + } + + let mut ts = TimeSlice::new(0i64, 0.0); + + ts.add_events( + vec![ + vec![vec![a], vec![b]], + vec![vec![c], vec![d]], + vec![vec![a], vec![c]], + ], + vec![vec![1.0, 0.0], vec![1.0, 0.0], vec![1.0, 0.0]], + vec![], + &agents, + ); + + assert_eq!(ts.color_groups.n_colors(), 2); + assert_eq!(ts.color_groups.groups[0], vec![0, 1]); + assert_eq!(ts.color_groups.groups[1], vec![2]); + + assert_eq!(ts.color_groups.color_range(0), 0..2); + assert_eq!(ts.color_groups.color_range(1), 2..3); + + // Events at positions 0 and 1 (color 0) must be disjoint — verify by + // checking that the agent sets of self.events[0] and self.events[1] do + // not include the agent at self.events[2]. + let agents_in_ev2: Vec = ts.events[2].iter_agents().collect(); + let agents_in_ev0: Vec = ts.events[0].iter_agents().collect(); + let agents_in_ev1: Vec = ts.events[1].iter_agents().collect(); + // ev0 and ev1 must be disjoint from each other (color-0 invariant). + assert!(agents_in_ev0.iter().all(|ag| !agents_in_ev1.contains(ag))); + // ev2 must share an agent with ev0 or ev1 (it needed its own color). + let ev2_overlaps_ev0 = agents_in_ev2.iter().any(|ag| agents_in_ev0.contains(ag)); + let ev2_overlaps_ev1 = agents_in_ev2.iter().any(|ag| agents_in_ev1.contains(ag)); + assert!(ev2_overlaps_ev0 || ev2_overlaps_ev1); + } }