1use serde::{Deserialize, Serialize};
27
28use crate::causal::CausalGraph;
29
30pub use eml_core::EmlEvent;
32pub use eml_core::{eml, eml_safe, softmax3};
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct CoherencePrediction {
41 pub lambda_2: f64,
43 pub fiedler_norm: f64,
45 pub uncertainty: f64,
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct GraphFeatures {
57 pub node_count: f64,
59 pub edge_count: f64,
61 pub avg_degree: f64,
63 pub max_degree: f64,
65 pub min_degree: f64,
67 pub density: f64,
69 pub component_count: f64,
71}
72
73impl GraphFeatures {
74 pub fn from_causal_graph(graph: &CausalGraph) -> Self {
76 let n = graph.node_count() as f64;
77 let m = graph.edge_count() as f64;
78
79 if n < 1.0 {
80 return Self {
81 node_count: 0.0,
82 edge_count: 0.0,
83 avg_degree: 0.0,
84 max_degree: 0.0,
85 min_degree: 0.0,
86 density: 0.0,
87 component_count: 0.0,
88 };
89 }
90
91 let ids = graph.node_ids();
92 let mut max_deg: usize = 0;
93 let mut min_deg: usize = usize::MAX;
94 for &id in &ids {
95 let d = graph.degree(id);
96 if d > max_deg {
97 max_deg = d;
98 }
99 if d < min_deg {
100 min_deg = d;
101 }
102 }
103 if ids.is_empty() {
104 min_deg = 0;
105 }
106
107 let avg_degree = if n > 0.0 { 2.0 * m / n } else { 0.0 };
108 let density = if n > 1.0 {
109 2.0 * m / (n * (n - 1.0))
110 } else {
111 0.0
112 };
113
114 let component_count = graph.component_count_fast() as f64;
116
117 Self {
118 node_count: n,
119 edge_count: m,
120 avg_degree,
121 max_degree: max_deg as f64,
122 min_degree: min_deg as f64,
123 density,
124 component_count,
125 }
126 }
127
128 fn normalized(&self) -> [f64; 7] {
130 [
131 self.node_count / 10000.0,
132 self.edge_count / 50000.0,
133 self.avg_degree / 100.0,
134 self.max_degree / 1000.0,
135 self.density,
136 self.component_count / 100.0,
137 self.min_degree / 50.0,
138 ]
139 }
140}
141
142impl eml_core::FeatureVector for GraphFeatures {
143 fn as_features(&self) -> Vec<f64> {
144 self.normalized().to_vec()
145 }
146
147 fn feature_count() -> usize {
148 7
149 }
150}
151
152#[allow(dead_code)] const PARAM_COUNT_V1: usize = 34;
159
160const PARAM_COUNT_V2: usize = 50;
162
163#[derive(Debug, Clone, Serialize, Deserialize)]
171pub struct EmlCoherenceModel {
172 inner: eml_core::EmlModel,
174 #[serde(skip)]
176 error_history: Vec<f64>,
177}
178
179impl Default for EmlCoherenceModel {
180 fn default() -> Self {
181 Self::new()
182 }
183}
184
185impl EmlCoherenceModel {
186 pub fn new() -> Self {
188 let mut inner = eml_core::EmlModel::new(4, 7, 3);
189 inner.set_model_name("coherence");
190 Self {
191 inner,
192 error_history: Vec::new(),
193 }
194 }
195
196 pub fn new_v1() -> Self {
198 let mut inner = eml_core::EmlModel::new(3, 7, 1);
199 inner.set_model_name("coherence_v1");
200 Self {
201 inner,
202 error_history: Vec::new(),
203 }
204 }
205
206 pub fn drain_events(&mut self) -> Vec<EmlEvent> {
208 self.inner.drain_events()
209 }
210
211 pub fn is_multi_head(&self) -> bool {
213 self.inner.param_count() == PARAM_COUNT_V2
214 }
215
216 pub fn is_trained(&self) -> bool {
218 self.inner.is_trained()
219 }
220
221 pub fn training_sample_count(&self) -> usize {
223 self.inner.training_sample_count()
224 }
225
226 pub fn mean_error(&self) -> f64 {
228 if self.error_history.is_empty() {
229 return 0.0;
230 }
231 self.error_history.iter().sum::<f64>() / self.error_history.len() as f64
232 }
233
234 pub fn predict(&self, features: &GraphFeatures) -> CoherencePrediction {
245 if !self.inner.is_trained() {
246 let lambda_2 = features.density * features.avg_degree;
249 return CoherencePrediction {
250 lambda_2,
251 fiedler_norm: lambda_2.sqrt().max(0.0),
252 uncertainty: lambda_2 * 0.5,
253 };
254 }
255
256 let inputs = features.normalized();
257 let heads = self.inner.predict(&inputs);
258
259 if self.inner.head_count() == 1 {
260 let lambda_2 = heads[0];
262 CoherencePrediction {
263 lambda_2,
264 fiedler_norm: lambda_2.sqrt().max(0.0),
265 uncertainty: lambda_2 * 0.5,
266 }
267 } else {
268 CoherencePrediction {
269 lambda_2: heads[0],
270 fiedler_norm: heads.get(1).copied().unwrap_or(0.0),
271 uncertainty: heads.get(2).copied().unwrap_or(0.0),
272 }
273 }
274 }
275
276 pub fn predict_lambda2(&self, features: &GraphFeatures) -> f64 {
281 self.predict(features).lambda_2
282 }
283
284 pub fn record(&mut self, features: GraphFeatures, lambda_2: f64) {
293 self.record_full(features, lambda_2, None, None);
294 }
295
296 pub fn record_full(
298 &mut self,
299 features: GraphFeatures,
300 lambda_2: f64,
301 fiedler_norm: Option<f64>,
302 uncertainty: Option<f64>,
303 ) {
304 let predicted = self.predict(&features);
306 self.error_history
307 .push((predicted.lambda_2 - lambda_2).abs());
308 if self.error_history.len() > 100 {
309 self.error_history.remove(0);
310 }
311
312 let inputs = features.normalized();
313 let targets = if self.inner.head_count() == 1 {
314 vec![Some(lambda_2)]
315 } else {
316 vec![Some(lambda_2), fiedler_norm, uncertainty]
317 };
318 self.inner.record(&inputs, &targets);
319 }
320
321 pub fn train(&mut self) -> bool {
328 self.inner.train()
329 }
330}
331
332impl CausalGraph {
337 pub fn coherence_fast(&self, model: &EmlCoherenceModel) -> CoherencePrediction {
343 let features = GraphFeatures::from_causal_graph(self);
344 model.predict(&features)
345 }
346}
347
348#[cfg(test)]
353mod tests {
354 use super::*;
355 use crate::causal::{CausalEdgeType, CausalGraph};
356
357 #[test]
360 fn eml_identity() {
361 let result = eml(0.0, 1.0);
363 assert!(
364 (result - 1.0).abs() < 1e-12,
365 "eml(0, 1) should be 1.0, got {result}"
366 );
367 }
368
369 #[test]
370 fn eml_exp_only() {
371 let result = eml(1.0, 1.0);
373 assert!(
374 (result - std::f64::consts::E).abs() < 1e-12,
375 "eml(1, 1) should be e, got {result}"
376 );
377 }
378
379 #[test]
380 fn eml_ln_only() {
381 let result = eml(0.0, std::f64::consts::E);
383 assert!(
384 result.abs() < 1e-12,
385 "eml(0, e) should be 0.0, got {result}"
386 );
387 }
388
389 #[test]
392 fn features_empty_graph() {
393 let g = CausalGraph::new();
394 let f = GraphFeatures::from_causal_graph(&g);
395 assert_eq!(f.node_count, 0.0);
396 assert_eq!(f.edge_count, 0.0);
397 assert_eq!(f.density, 0.0);
398 assert_eq!(f.component_count, 0.0);
399 }
400
401 #[test]
402 fn features_triangle() {
403 let g = CausalGraph::new();
404 let a = g.add_node("A".into(), serde_json::json!({}));
405 let b = g.add_node("B".into(), serde_json::json!({}));
406 let c = g.add_node("C".into(), serde_json::json!({}));
407 g.link(a, b, CausalEdgeType::Causes, 1.0, 0, 0);
408 g.link(b, c, CausalEdgeType::Causes, 1.0, 0, 0);
409 g.link(c, a, CausalEdgeType::Causes, 1.0, 0, 0);
410
411 let f = GraphFeatures::from_causal_graph(&g);
412 assert_eq!(f.node_count, 3.0);
413 assert_eq!(f.edge_count, 3.0);
414 assert!((f.avg_degree - 2.0).abs() < 1e-9);
415 assert_eq!(f.component_count, 1.0);
416 assert!((f.density - 1.0).abs() < 1e-9);
418 }
419
420 #[test]
421 fn features_disconnected() {
422 let g = CausalGraph::new();
423 let _a = g.add_node("A".into(), serde_json::json!({}));
424 let _b = g.add_node("B".into(), serde_json::json!({}));
425
426 let f = GraphFeatures::from_causal_graph(&g);
427 assert_eq!(f.node_count, 2.0);
428 assert_eq!(f.edge_count, 0.0);
429 assert_eq!(f.component_count, 2.0);
430 assert_eq!(f.min_degree, 0.0);
431 assert_eq!(f.max_degree, 0.0);
432 }
433
434 #[test]
437 fn predict_untrained_fallback() {
438 let model = EmlCoherenceModel::new();
439 assert!(!model.is_trained());
440
441 let features = GraphFeatures {
442 node_count: 10.0,
443 edge_count: 20.0,
444 avg_degree: 4.0,
445 max_degree: 6.0,
446 min_degree: 2.0,
447 density: 0.444,
448 component_count: 1.0,
449 };
450 let result = model.predict(&features);
451 let expected = 0.444 * 4.0;
453 assert!(
454 (result.lambda_2 - expected).abs() < 1e-9,
455 "untrained fallback: expected {expected}, got {}",
456 result.lambda_2
457 );
458 }
459
460 #[test]
461 fn predict_untrained_returns_multi_head() {
462 let model = EmlCoherenceModel::new();
463 let features = GraphFeatures {
464 node_count: 10.0,
465 edge_count: 20.0,
466 avg_degree: 4.0,
467 max_degree: 6.0,
468 min_degree: 2.0,
469 density: 0.444,
470 component_count: 1.0,
471 };
472 let pred = model.predict(&features);
473 assert!(pred.lambda_2 >= 0.0);
475 assert!(pred.fiedler_norm >= 0.0);
476 assert!(pred.uncertainty >= 0.0);
477 }
478
479 #[test]
480 fn predict_lambda2_convenience() {
481 let model = EmlCoherenceModel::new();
482 let features = GraphFeatures {
483 node_count: 10.0,
484 edge_count: 20.0,
485 avg_degree: 4.0,
486 max_degree: 6.0,
487 min_degree: 2.0,
488 density: 0.444,
489 component_count: 1.0,
490 };
491 let lambda2 = model.predict_lambda2(&features);
492 let full = model.predict(&features);
493 assert!(
494 (lambda2 - full.lambda_2).abs() < 1e-12,
495 "predict_lambda2 should match predict().lambda_2"
496 );
497 }
498
499 #[test]
502 fn backward_compat_v1_model() {
503 let mut model = EmlCoherenceModel::new_v1();
504 assert!(!model.is_multi_head());
505
506 let features = GraphFeatures {
507 node_count: 10.0,
508 edge_count: 20.0,
509 avg_degree: 4.0,
510 max_degree: 6.0,
511 min_degree: 2.0,
512 density: 0.444,
513 component_count: 1.0,
514 };
515
516 let pred = model.predict(&features);
518 let expected = 0.444 * 4.0;
519 assert!(
520 (pred.lambda_2 - expected).abs() < 1e-9,
521 "v1 untrained fallback should match"
522 );
523 assert!(pred.fiedler_norm >= 0.0);
525 assert!(pred.uncertainty >= 0.0);
526
527 for i in 0..60 {
529 let f = GraphFeatures {
530 node_count: (i + 3) as f64,
531 edge_count: (i + 2) as f64,
532 avg_degree: 2.0,
533 max_degree: 3.0,
534 min_degree: 1.0,
535 density: 0.5,
536 component_count: 1.0,
537 };
538 model.record(f, 1.0);
539 }
540 let _ = model.train();
542 }
543
544 #[test]
547 fn new_model_is_multi_head() {
548 let model = EmlCoherenceModel::new();
549 assert!(model.is_multi_head());
550 }
551
552 #[test]
553 fn uncertainty_is_non_negative() {
554 let model = EmlCoherenceModel::new();
555 for n in [3.0, 10.0, 50.0, 100.0] {
557 for d in [0.1, 0.5, 1.0] {
558 let e = n * (n - 1.0) * d / 2.0;
559 let features = GraphFeatures {
560 node_count: n,
561 edge_count: e,
562 avg_degree: (n - 1.0) * d,
563 max_degree: (n - 1.0) * d * 1.5,
564 min_degree: ((n - 1.0) * d * 0.5).max(0.0),
565 density: d,
566 component_count: 1.0,
567 };
568 let pred = model.predict(&features);
569 assert!(
570 pred.uncertainty >= 0.0,
571 "uncertainty must be non-negative for n={n}, d={d}: got {}",
572 pred.uncertainty
573 );
574 }
575 }
576 }
577
578 #[test]
581 fn record_increments_count() {
582 let mut model = EmlCoherenceModel::new();
583 assert_eq!(model.training_sample_count(), 0);
584
585 let f = GraphFeatures {
586 node_count: 5.0,
587 edge_count: 4.0,
588 avg_degree: 1.6,
589 max_degree: 2.0,
590 min_degree: 1.0,
591 density: 0.4,
592 component_count: 1.0,
593 };
594 model.record(f, 0.5);
595 assert_eq!(model.training_sample_count(), 1);
596 }
597
598 #[test]
599 fn record_full_stores_targets() {
600 let mut model = EmlCoherenceModel::new();
601 let f = GraphFeatures {
602 node_count: 5.0,
603 edge_count: 4.0,
604 avg_degree: 1.6,
605 max_degree: 2.0,
606 min_degree: 1.0,
607 density: 0.4,
608 component_count: 1.0,
609 };
610 model.record_full(f, 0.5, Some(1.2), Some(0.3));
611 assert_eq!(model.training_sample_count(), 1);
612 }
613
614 #[test]
615 fn train_insufficient_data_returns_false() {
616 let mut model = EmlCoherenceModel::new();
617 for i in 0..10 {
619 let f = GraphFeatures {
620 node_count: i as f64,
621 edge_count: i as f64,
622 avg_degree: 2.0,
623 max_degree: 3.0,
624 min_degree: 1.0,
625 density: 0.5,
626 component_count: 1.0,
627 };
628 model.record(f, 1.0);
629 }
630 assert!(!model.train());
631 assert!(!model.is_trained());
632 }
633
634 #[test]
637 fn convergence_on_known_graphs() {
638 let mut model = EmlCoherenceModel::new();
639
640 for n in 3..30 {
642 let nf = n as f64;
643 let e = nf * (nf - 1.0) / 2.0;
644 let lambda_2 = nf;
645 let features = GraphFeatures {
646 node_count: nf,
647 edge_count: e,
648 avg_degree: nf - 1.0,
649 max_degree: nf - 1.0,
650 min_degree: nf - 1.0,
651 density: 1.0,
652 component_count: 1.0,
653 };
654 model.record(features, lambda_2);
655 }
656
657 for n in 3..30 {
659 let nf = n as f64;
660 let features = GraphFeatures {
661 node_count: nf,
662 edge_count: nf - 1.0,
663 avg_degree: 2.0 * (nf - 1.0) / nf,
664 max_degree: nf - 1.0,
665 min_degree: 1.0,
666 density: 2.0 * (nf - 1.0) / (nf * (nf - 1.0)),
667 component_count: 1.0,
668 };
669 model.record(features, 1.0);
670 }
671
672 for n in 3..30 {
674 let nf = n as f64;
675 let lambda_2 = 2.0 * (1.0 - (2.0 * std::f64::consts::PI / nf).cos());
676 let features = GraphFeatures {
677 node_count: nf,
678 edge_count: nf,
679 avg_degree: 2.0,
680 max_degree: 2.0,
681 min_degree: 2.0,
682 density: 2.0 * nf / (nf * (nf - 1.0)),
683 component_count: 1.0,
684 };
685 model.record(features, lambda_2);
686 }
687
688 for n in 3..30 {
690 let nf = n as f64;
691 let lambda_2 = 2.0 * (1.0 - (std::f64::consts::PI / nf).cos());
692 let features = GraphFeatures {
693 node_count: nf,
694 edge_count: nf - 1.0,
695 avg_degree: 2.0 * (nf - 1.0) / nf,
696 max_degree: 2.0,
697 min_degree: 1.0,
698 density: 2.0 * (nf - 1.0) / (nf * (nf - 1.0)),
699 component_count: 1.0,
700 };
701 model.record(features, lambda_2);
702 }
703
704 for n in [20, 50, 100, 200] {
706 for &p in &[0.1, 0.2, 0.3, 0.5, 0.7] {
707 let nf = n as f64;
708 let e = nf * (nf - 1.0) * p / 2.0;
709 let avg_deg = (nf - 1.0) * p;
710 let lambda_2 = (nf * p - 2.0 * (nf * p * (1.0 - p)).sqrt()).max(0.0);
711 let features = GraphFeatures {
712 node_count: nf,
713 edge_count: e,
714 avg_degree: avg_deg,
715 max_degree: avg_deg * 1.5,
716 min_degree: (avg_deg * 0.5).max(0.0),
717 density: p,
718 component_count: 1.0,
719 };
720 model.record(features, lambda_2);
721 }
722 }
723
724 assert!(
725 model.training_sample_count() >= 50,
726 "should have enough training data: {}",
727 model.training_sample_count()
728 );
729
730 let _converged = model.train();
732
733 let k5 = GraphFeatures {
735 node_count: 5.0,
736 edge_count: 10.0,
737 avg_degree: 4.0,
738 max_degree: 4.0,
739 min_degree: 4.0,
740 density: 1.0,
741 component_count: 1.0,
742 };
743 let pred = model.predict(&k5);
744 assert!(pred.lambda_2.is_finite(), "prediction should be finite");
745 }
746
747 #[test]
748 fn depth4_convergence_with_full_targets() {
749 let mut model = EmlCoherenceModel::new();
750 assert!(model.is_multi_head());
751
752 for n in 3..53 {
754 let nf = n as f64;
755 let e = nf * (nf - 1.0) / 2.0;
756 let lambda_2 = nf; let fiedler_norm = (nf - 1.0).sqrt();
758 let uncertainty = 0.1 * lambda_2;
759 let features = GraphFeatures {
760 node_count: nf,
761 edge_count: e,
762 avg_degree: nf - 1.0,
763 max_degree: nf - 1.0,
764 min_degree: nf - 1.0,
765 density: 1.0,
766 component_count: 1.0,
767 };
768 model.record_full(features, lambda_2, Some(fiedler_norm), Some(uncertainty));
769 }
770
771 for n in 3..53 {
772 let nf = n as f64;
773 let lambda_2 = 2.0 * (1.0 - (2.0 * std::f64::consts::PI / nf).cos());
774 let fiedler_norm = lambda_2.sqrt();
775 let uncertainty = 0.05 * lambda_2;
776 let features = GraphFeatures {
777 node_count: nf,
778 edge_count: nf,
779 avg_degree: 2.0,
780 max_degree: 2.0,
781 min_degree: 2.0,
782 density: 2.0 * nf / (nf * (nf - 1.0)),
783 component_count: 1.0,
784 };
785 model.record_full(features, lambda_2, Some(fiedler_norm), Some(uncertainty));
786 }
787
788 assert!(model.training_sample_count() >= 100);
789
790 let _ = model.train();
792 let pred = model.predict(&GraphFeatures {
793 node_count: 10.0,
794 edge_count: 45.0,
795 avg_degree: 9.0,
796 max_degree: 9.0,
797 min_degree: 9.0,
798 density: 1.0,
799 component_count: 1.0,
800 });
801 assert!(pred.lambda_2.is_finite());
802 }
803
804 #[test]
807 fn coherence_fast_on_triangle() {
808 let g = CausalGraph::new();
809 let a = g.add_node("A".into(), serde_json::json!({}));
810 let b = g.add_node("B".into(), serde_json::json!({}));
811 let c = g.add_node("C".into(), serde_json::json!({}));
812 g.link(a, b, CausalEdgeType::Causes, 1.0, 0, 0);
813 g.link(b, c, CausalEdgeType::Causes, 1.0, 0, 0);
814 g.link(c, a, CausalEdgeType::Causes, 1.0, 0, 0);
815
816 let model = EmlCoherenceModel::new();
817 let fast = g.coherence_fast(&model);
818 assert!(
820 (fast.lambda_2 - 2.0).abs() < 1e-9,
821 "coherence_fast untrained triangle: expected 2.0, got {}",
822 fast.lambda_2
823 );
824 assert!(fast.fiedler_norm >= 0.0);
826 assert!(fast.uncertainty >= 0.0);
827 }
828
829 #[test]
830 fn coherence_fast_empty() {
831 let g = CausalGraph::new();
832 let model = EmlCoherenceModel::new();
833 let fast = g.coherence_fast(&model);
834 assert!(
835 fast.lambda_2.abs() < 1e-12,
836 "coherence_fast on empty graph should be 0"
837 );
838 }
839
840 #[test]
843 fn softmax3_sums_to_one() {
844 let (a, b, c) = softmax3(1.0, 2.0, 3.0);
845 let sum = a + b + c;
846 assert!(
847 (sum - 1.0).abs() < 1e-12,
848 "softmax3 should sum to 1.0, got {sum}"
849 );
850 }
851
852 #[test]
853 fn softmax3_equal_inputs() {
854 let (a, b, c) = softmax3(0.0, 0.0, 0.0);
855 assert!((a - 1.0 / 3.0).abs() < 1e-12);
856 assert!((b - 1.0 / 3.0).abs() < 1e-12);
857 assert!((c - 1.0 / 3.0).abs() < 1e-12);
858 }
859
860 #[test]
861 fn eml_safe_does_not_panic() {
862 let _ = eml_safe(100.0, 0.0);
864 let _ = eml_safe(-100.0, -5.0);
865 let _ = eml_safe(0.0, f64::MIN_POSITIVE);
866 let _ = eml_safe(f64::NAN, 1.0); }
868
869 #[test]
870 fn error_history_tracks_drift() {
871 let mut model = EmlCoherenceModel::new();
872 let f = GraphFeatures {
873 node_count: 5.0,
874 edge_count: 5.0,
875 avg_degree: 2.0,
876 max_degree: 2.0,
877 min_degree: 2.0,
878 density: 0.5,
879 component_count: 1.0,
880 };
881
882 model.record(f.clone(), 1.0);
883 model.record(f.clone(), 2.0);
884 assert_eq!(model.error_history.len(), 2);
885 assert!(model.mean_error() >= 0.0);
886 }
887}