1use petgraph::EdgeType;
4use petgraph::graph::NodeIndex;
5use petgraph::stable_graph::StableGraph;
6
7#[derive(Debug, Clone, Copy)]
9pub struct SokmConfig {
10 pub decay_factor: f64,
12 pub delta: f64,
14 pub min_weight: f64,
16 pub formula: StrengthFormula,
18}
19
20impl Default for SokmConfig {
21 fn default() -> Self {
22 Self {
23 decay_factor: 0.95,
24 delta: 0.02,
25 min_weight: 0.001,
26 formula: StrengthFormula::default(),
27 }
28 }
29}
30
31#[derive(Debug, Clone, Copy, Default)]
33pub enum StrengthFormula {
34 #[default]
36 Product,
37 Min,
39 Average,
41}
42
43#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
45pub struct HebbianReport {
46 pub decayed: usize,
48 pub strengthened: usize,
50 pub pruned: usize,
52}
53
54pub fn decay<N, Ty: EdgeType>(graph: &mut StableGraph<N, f64, Ty>, factor: f64) -> usize {
71 let indices: Vec<_> = graph.edge_indices().collect();
72 for &idx in &indices {
73 if let Some(w) = graph.edge_weight_mut(idx) {
74 *w *= factor;
75 }
76 }
77 indices.len()
78}
79
80pub fn strengthen<N, Ty: EdgeType>(
108 graph: &mut StableGraph<N, f64, Ty>,
109 activated: &[(NodeIndex, f64)],
110 config: &SokmConfig,
111) -> usize {
112 let mut count = 0;
113 for i in 0..activated.len() {
114 let j_start = if Ty::is_directed() { 0 } else { i + 1 };
115 for j in j_start..activated.len() {
116 if i == j {
117 continue;
118 }
119 let (na, sa) = activated[i];
120 let (nb, sb) = activated[j];
121 let increment = match config.formula {
122 StrengthFormula::Product => config.delta * sa * sb,
123 StrengthFormula::Min => config.delta * sa.min(sb),
124 StrengthFormula::Average => config.delta * (sa + sb) / 2.0,
125 };
126 if let Some(idx) = graph.find_edge(na, nb) {
127 if let Some(w) = graph.edge_weight_mut(idx) {
128 *w += increment;
129 }
130 } else {
131 graph.add_edge(na, nb, increment);
132 }
133 count += 1;
134 }
135 }
136 count
137}
138
139pub fn prune<N, Ty: EdgeType>(graph: &mut StableGraph<N, f64, Ty>, threshold: f64) -> usize {
157 let to_remove: Vec<_> = graph
158 .edge_indices()
159 .filter(|&idx| graph.edge_weight(idx).is_some_and(|&w| w < threshold))
160 .collect();
161 let count = to_remove.len();
162 for idx in to_remove {
163 graph.remove_edge(idx);
164 }
165 count
166}
167
168pub fn sokm_tick<N, Ty: EdgeType>(
186 graph: &mut StableGraph<N, f64, Ty>,
187 activated: &[(NodeIndex, f64)],
188 config: &SokmConfig,
189) -> HebbianReport {
190 let decayed = decay(graph, config.decay_factor);
191 let strengthened = strengthen(graph, activated, config);
192 let pruned = prune(graph, config.min_weight);
193 HebbianReport {
194 decayed,
195 strengthened,
196 pruned,
197 }
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203 use petgraph::stable_graph::{StableDiGraph, StableUnGraph};
204
205 fn make_graph() -> StableDiGraph<(), f64> {
206 let mut g = StableDiGraph::new();
207 let a = g.add_node(());
208 let b = g.add_node(());
209 let c = g.add_node(());
210 g.add_edge(a, b, 0.5);
211 g.add_edge(b, c, 0.3);
212 g.add_edge(a, c, 0.01);
213 g
214 }
215
216 #[test]
217 fn decay_multiplies_all_weights() {
218 let mut g = make_graph();
219 assert_eq!(decay(&mut g, 0.5), 3);
220 assert_eq!(*g.edge_weight(0.into()).unwrap(), 0.25);
221 assert_eq!(*g.edge_weight(1.into()).unwrap(), 0.15);
222 }
223
224 #[test]
225 fn decay_repeated_converges() {
226 let mut g = StableDiGraph::<(), f64>::new();
227 let a = g.add_node(());
228 let b = g.add_node(());
229 g.add_edge(a, b, 1.0);
230
231 for _ in 0..10 {
232 decay(&mut g, 0.95);
233 }
234 let expected = 0.95f64.powi(10);
235 assert!((g.edge_weight(0.into()).unwrap() - expected).abs() < 1e-10);
236 }
237
238 #[test]
239 fn prune_removes_below_threshold() {
240 let mut g = make_graph();
241 assert_eq!(prune(&mut g, 0.1), 1);
242 assert_eq!(g.edge_count(), 2);
243 }
244
245 #[test]
246 fn prune_retains_above_threshold() {
247 let mut g = make_graph();
248 assert_eq!(prune(&mut g, 0.001), 0);
249 assert_eq!(g.edge_count(), 3);
250 }
251
252 #[test]
253 fn strengthen_creates_missing_edges() {
254 let mut g = StableDiGraph::<(), f64>::new();
255 let a = g.add_node(());
256 let b = g.add_node(());
257
258 let config = SokmConfig::default();
259 strengthen(&mut g, &[(a, 1.0), (b, 1.0)], &config);
260 assert_eq!(g.edge_count(), 2); }
262
263 #[test]
264 fn strengthen_higher_scores_get_more() {
265 let mut g1 = StableDiGraph::<(), f64>::new();
266 let a1 = g1.add_node(());
267 let b1 = g1.add_node(());
268
269 let mut g2 = StableDiGraph::<(), f64>::new();
270 let a2 = g2.add_node(());
271 let b2 = g2.add_node(());
272
273 let config = SokmConfig::default();
274 strengthen(&mut g1, &[(a1, 0.5), (b1, 0.5)], &config);
275 strengthen(&mut g2, &[(a2, 1.0), (b2, 1.0)], &config);
276
277 let w1 = *g1.edge_weight(g1.find_edge(a1, b1).unwrap()).unwrap();
278 let w2 = *g2.edge_weight(g2.find_edge(a2, b2).unwrap()).unwrap();
279 assert!(w2 > w1);
280 }
281
282 #[test]
283 fn strengthen_non_activated_pairs_unchanged() {
284 let mut g = StableDiGraph::<(), f64>::new();
285 let a = g.add_node(());
286 let _b = g.add_node(());
287 let c = g.add_node(());
288 g.add_edge(a, c, 0.5);
289
290 let config = SokmConfig::default();
291 strengthen(&mut g, &[(a, 1.0)], &config);
292 assert_eq!(*g.edge_weight(0.into()).unwrap(), 0.5);
293 }
294
295 #[test]
296 fn strengthen_increments_existing_edge() {
297 let mut g = StableDiGraph::<(), f64>::new();
298 let a = g.add_node(());
299 let b = g.add_node(());
300 g.add_edge(a, b, 0.5);
301
302 let config = SokmConfig::default();
303 strengthen(&mut g, &[(a, 1.0), (b, 1.0)], &config);
304 assert!((g.edge_weight(0.into()).unwrap() - 0.52).abs() < 1e-10);
305 }
306
307 #[test]
308 fn sokm_tick_full_cycle() {
309 let mut g = StableDiGraph::<(), f64>::new();
310 let a = g.add_node(());
311 let b = g.add_node(());
312 g.add_edge(a, b, 0.0005);
313
314 let config = SokmConfig::default();
315 let report = sokm_tick(&mut g, &[(a, 1.0), (b, 1.0)], &config);
316 assert_eq!(report.decayed, 1);
317 assert_eq!(report.strengthened, 2);
318 assert_eq!(report.pruned, 0);
319 }
320
321 #[test]
322 fn sokm_tick_prunes_weak_edge() {
323 let mut g = StableDiGraph::<(), f64>::new();
324 let a = g.add_node(());
325 let b = g.add_node(());
326 let c = g.add_node(());
327 g.add_edge(a, c, 0.0005);
328
329 let config = SokmConfig::default();
330 let report = sokm_tick(&mut g, &[(a, 1.0), (b, 1.0)], &config);
331 assert_eq!(report.pruned, 1);
332 }
333
334 #[test]
335 fn strength_formula_min() {
336 let mut g = StableDiGraph::<(), f64>::new();
337 let a = g.add_node(());
338 let b = g.add_node(());
339
340 let config = SokmConfig {
341 formula: StrengthFormula::Min,
342 ..SokmConfig::default()
343 };
344 strengthen(&mut g, &[(a, 0.8), (b, 0.3)], &config);
345 assert!((g.edge_weight(0.into()).unwrap() - 0.006).abs() < 1e-10);
346 }
347
348 #[test]
349 fn strength_formula_average() {
350 let mut g = StableDiGraph::<(), f64>::new();
351 let a = g.add_node(());
352 let b = g.add_node(());
353
354 let config = SokmConfig {
355 formula: StrengthFormula::Average,
356 ..SokmConfig::default()
357 };
358 strengthen(&mut g, &[(a, 0.8), (b, 0.4)], &config);
359 assert!((g.edge_weight(0.into()).unwrap() - 0.012).abs() < 1e-10);
360 }
361
362 #[test]
363 fn undirected_strengthen_single_edge_per_pair() {
364 let mut g = StableUnGraph::<(), f64>::with_capacity(0, 0);
365 let a = g.add_node(());
366 let b = g.add_node(());
367
368 let config = SokmConfig::default();
369 let count = strengthen(&mut g, &[(a, 1.0), (b, 1.0)], &config);
370 assert_eq!(count, 1);
372 assert_eq!(g.edge_count(), 1);
373 assert!((g.edge_weight(0.into()).unwrap() - 0.02).abs() < 1e-10);
374 }
375
376 #[test]
377 fn undirected_sokm_tick() {
378 let mut g = StableUnGraph::<(), f64>::with_capacity(0, 0);
379 let a = g.add_node(());
380 let b = g.add_node(());
381 let c = g.add_node(());
382 g.add_edge(a, b, 0.5);
383 g.add_edge(b, c, 0.0005);
384
385 let config = SokmConfig::default();
386 let report = sokm_tick(&mut g, &[(a, 1.0), (b, 0.9)], &config);
387 assert_eq!(report.decayed, 2);
388 assert_eq!(report.strengthened, 1);
389 assert_eq!(report.pruned, 1);
391 }
392}