swarm_engine_core/exploration/selection/
any.rs1use std::fmt::Debug;
7
8use super::{Fifo, Greedy, SelectionKind, SelectionLogic, Thompson, Ucb1};
9use crate::exploration::map::{GraphMap, MapNodeId, MapState};
10use crate::exploration::mutation::ActionExtractor;
11use crate::learn::LearnedProvider;
12use crate::online_stats::SwarmStats;
13
14#[derive(Debug, Clone)]
19pub enum AnySelection {
20 Fifo(Fifo),
21 Ucb1(Ucb1),
22 Greedy(Greedy),
23 Thompson(Thompson),
24}
25
26impl Default for AnySelection {
27 fn default() -> Self {
28 Self::Fifo(Fifo::new())
29 }
30}
31
32impl AnySelection {
33 pub fn from_kind(kind: SelectionKind, ucb1_c: f64) -> Self {
39 match kind {
40 SelectionKind::Fifo => Self::Fifo(Fifo::new()),
41 SelectionKind::Ucb1 => Self::Ucb1(Ucb1::new(ucb1_c)),
42 SelectionKind::Greedy => Self::Greedy(Greedy::new()),
43 SelectionKind::Thompson => Self::Thompson(Thompson::new()),
44 }
45 }
46
47 pub fn kind(&self) -> SelectionKind {
49 match self {
50 Self::Fifo(_) => SelectionKind::Fifo,
51 Self::Ucb1(_) => SelectionKind::Ucb1,
52 Self::Greedy(_) => SelectionKind::Greedy,
53 Self::Thompson(_) => SelectionKind::Thompson,
54 }
55 }
56
57 pub fn selection_name(&self) -> &str {
59 match self {
60 Self::Fifo(_) => "FIFO",
61 Self::Ucb1(_) => "UCB1",
62 Self::Greedy(_) => "Greedy",
63 Self::Thompson(_) => "Thompson",
64 }
65 }
66
67 pub fn compute_score(
69 &self,
70 action: &str,
71 target: Option<&str>,
72 stats: &SwarmStats,
73 provider: &dyn LearnedProvider,
74 ) -> f64 {
75 match self {
76 Self::Fifo(_) => 0.0,
77 Self::Ucb1(s) => s.compute_score(stats, action, target, provider),
78 Self::Greedy(s) => s.compute_score(action, target, provider),
79 Self::Thompson(s) => s.compute_score(stats, action, target, provider),
80 }
81 }
82}
83
84impl<N, E, S> SelectionLogic<N, E, S> for AnySelection
85where
86 N: Debug + Clone + ActionExtractor,
87 E: Debug + Clone,
88 S: MapState,
89{
90 fn next(
91 &self,
92 map: &GraphMap<N, E, S>,
93 stats: &SwarmStats,
94 provider: &dyn LearnedProvider,
95 ) -> Option<MapNodeId> {
96 match self {
97 Self::Fifo(s) => s.next(map, stats, provider),
98 Self::Ucb1(s) => s.next(map, stats, provider),
99 Self::Greedy(s) => s.next(map, stats, provider),
100 Self::Thompson(s) => s.next(map, stats, provider),
101 }
102 }
103
104 fn select(
105 &self,
106 map: &GraphMap<N, E, S>,
107 count: usize,
108 stats: &SwarmStats,
109 provider: &dyn LearnedProvider,
110 ) -> Vec<MapNodeId> {
111 match self {
112 Self::Fifo(s) => s.select(map, count, stats, provider),
113 Self::Ucb1(s) => s.select(map, count, stats, provider),
114 Self::Greedy(s) => s.select(map, count, stats, provider),
115 Self::Thompson(s) => s.select(map, count, stats, provider),
116 }
117 }
118
119 fn score(
120 &self,
121 action: &str,
122 target: Option<&str>,
123 stats: &SwarmStats,
124 provider: &dyn LearnedProvider,
125 ) -> f64 {
126 self.compute_score(action, target, stats, provider)
127 }
128
129 fn name(&self) -> &str {
130 self.selection_name()
131 }
132}
133
134#[cfg(test)]
135mod tests {
136 use super::*;
137 use crate::learn::NullProvider;
138
139 #[test]
140 fn test_any_selection_from_kind() {
141 let fifo = AnySelection::from_kind(SelectionKind::Fifo, 1.0);
142 assert_eq!(fifo.kind(), SelectionKind::Fifo);
143 assert_eq!(fifo.selection_name(), "FIFO");
144
145 let ucb1 = AnySelection::from_kind(SelectionKind::Ucb1, 1.41);
146 assert_eq!(ucb1.kind(), SelectionKind::Ucb1);
147 assert_eq!(ucb1.selection_name(), "UCB1");
148
149 let greedy = AnySelection::from_kind(SelectionKind::Greedy, 1.0);
150 assert_eq!(greedy.kind(), SelectionKind::Greedy);
151 assert_eq!(greedy.selection_name(), "Greedy");
152
153 let thompson = AnySelection::from_kind(SelectionKind::Thompson, 1.0);
154 assert_eq!(thompson.kind(), SelectionKind::Thompson);
155 assert_eq!(thompson.selection_name(), "Thompson");
156 }
157
158 #[test]
159 fn test_any_selection_score() {
160 let stats = SwarmStats::new();
161 let provider = NullProvider;
162
163 let fifo = AnySelection::from_kind(SelectionKind::Fifo, 1.0);
164 assert_eq!(fifo.compute_score("grep", None, &stats, &provider), 0.0);
165
166 let ucb1 = AnySelection::from_kind(SelectionKind::Ucb1, 1.41);
167 assert!(ucb1
169 .compute_score("grep", None, &stats, &provider)
170 .is_infinite());
171
172 let greedy = AnySelection::from_kind(SelectionKind::Greedy, 1.0);
173 assert_eq!(greedy.compute_score("grep", None, &stats, &provider), 0.5); let thompson = AnySelection::from_kind(SelectionKind::Thompson, 1.0);
176 assert_eq!(thompson.compute_score("grep", None, &stats, &provider), 0.5);
178 }
179
180 #[test]
181 fn test_any_selection_default() {
182 let any = AnySelection::default();
183 assert_eq!(any.kind(), SelectionKind::Fifo);
184 }
185}