petgraph_live/hebbian/
oja.rs1use petgraph::EdgeType;
7use petgraph::graph::NodeIndex;
8use petgraph::stable_graph::StableGraph;
9
10#[derive(Debug, Clone, Copy)]
12pub struct OjaConfig {
13 pub learning_rate: f64,
15}
16
17impl Default for OjaConfig {
18 fn default() -> Self {
19 Self {
20 learning_rate: 0.01,
21 }
22 }
23}
24
25pub fn oja_update<N, Ty: EdgeType>(
51 graph: &mut StableGraph<N, f64, Ty>,
52 pre_activations: &[(NodeIndex, f64)],
53 post_activations: &[(NodeIndex, f64)],
54 config: &OjaConfig,
55) -> usize {
56 let mut modified = 0;
57 for &(post_node, y) in post_activations {
58 for &(pre_node, x) in pre_activations {
59 if pre_node == post_node {
60 continue;
61 }
62 if let Some(idx) = graph.find_edge(pre_node, post_node)
63 && let Some(w) = graph.edge_weight_mut(idx)
64 {
65 *w += config.learning_rate * y * (x - *w * y);
67 modified += 1;
68 }
69 }
70 }
71 modified
72}
73
74#[cfg(test)]
75mod tests {
76 use super::*;
77 use petgraph::stable_graph::StableDiGraph;
78
79 #[test]
80 fn oja_converges_toward_normalization() {
81 let mut g = StableDiGraph::<(), f64>::new();
82 let a = g.add_node(());
83 let b = g.add_node(());
84 g.add_edge(a, b, 0.5);
85
86 let config = OjaConfig { learning_rate: 0.1 };
87 for _ in 0..100 {
89 oja_update(&mut g, &[(a, 1.0)], &[(b, 1.0)], &config);
90 }
91 let w = *g.edge_weight(0.into()).unwrap();
93 assert!((w - 1.0).abs() < 0.01);
94 }
95
96 #[test]
97 fn oja_no_edge_no_effect() {
98 let mut g = StableDiGraph::<(), f64>::new();
99 let a = g.add_node(());
100 let b = g.add_node(());
101
102 let modified = oja_update(&mut g, &[(a, 1.0)], &[(b, 1.0)], &OjaConfig::default());
103 assert_eq!(modified, 0);
104 assert_eq!(g.edge_count(), 0);
105 }
106
107 #[test]
108 fn oja_self_loop_skipped() {
109 let mut g = StableDiGraph::<(), f64>::new();
110 let a = g.add_node(());
111 g.add_edge(a, a, 0.5);
112
113 let modified = oja_update(&mut g, &[(a, 1.0)], &[(a, 1.0)], &OjaConfig::default());
114 assert_eq!(modified, 0);
115 assert_eq!(*g.edge_weight(0.into()).unwrap(), 0.5);
116 }
117
118 #[test]
119 fn oja_high_weight_decreases() {
120 let mut g = StableDiGraph::<(), f64>::new();
121 let a = g.add_node(());
122 let b = g.add_node(());
123 g.add_edge(a, b, 2.0); let config = OjaConfig { learning_rate: 0.1 };
126 oja_update(&mut g, &[(a, 1.0)], &[(b, 1.0)], &config);
127 assert!(*g.edge_weight(0.into()).unwrap() < 2.0);
129 }
130
131 #[test]
132 fn oja_low_weight_increases() {
133 let mut g = StableDiGraph::<(), f64>::new();
134 let a = g.add_node(());
135 let b = g.add_node(());
136 g.add_edge(a, b, 0.1); let config = OjaConfig { learning_rate: 0.1 };
139 oja_update(&mut g, &[(a, 1.0)], &[(b, 1.0)], &config);
140 assert!(*g.edge_weight(0.into()).unwrap() > 0.1);
141 }
142}