1use scirs2_core::numeric::Float;
13use std::cmp::Ordering;
14
15pub fn total_order<A: Float>(a: &A, b: &A) -> Ordering {
24 crate::utils::total_order(a, b)
25}
26
27pub fn sort_ascending<A: Float>(values: &mut [A]) {
29 values.sort_by(total_order);
30}
31
32pub fn mean<A: Float>(values: &[A]) -> Option<A> {
34 if values.is_empty() {
35 return None;
36 }
37 let sum = values.iter().fold(A::zero(), |acc, &v| acc + v);
38 A::from(values.len()).map(|n| sum / n)
39}
40
41pub fn population_variance<A: Float>(values: &[A]) -> Option<A> {
43 let mu = mean(values)?;
44 let n = A::from(values.len())?;
45 let ss = values.iter().fold(A::zero(), |acc, &v| {
46 let d = v - mu;
47 acc + d * d
48 });
49 Some(ss / n)
50}
51
52pub fn sample_std_dev<A: Float>(values: &[A]) -> Option<A> {
54 if values.len() < 2 {
55 return None;
56 }
57 let mu = mean(values)?;
58 let denom = A::from(values.len() - 1)?;
59 let ss = values.iter().fold(A::zero(), |acc, &v| {
60 let d = v - mu;
61 acc + d * d
62 });
63 Some((ss / denom).sqrt())
64}
65
66pub fn quantile_in_place<A: Float>(values: &mut [A], p: f64) -> Option<A> {
69 let n = values.len();
70 if n == 0 {
71 return None;
72 }
73 let p = p.clamp(0.0, 1.0);
74 let rank = (p * n as f64).ceil();
75 let index = if rank < 1.0 {
76 0
77 } else {
78 ((rank as usize) - 1).min(n - 1)
79 };
80 let (_, nth, _) = values.select_nth_unstable_by(index, total_order);
81 Some(*nth)
82}
83
84pub fn median_in_place<A: Float>(values: &mut [A]) -> Option<A> {
87 let n = values.len();
88 if n == 0 {
89 return None;
90 }
91 if n % 2 == 1 {
92 let (_, nth, _) = values.select_nth_unstable_by(n / 2, total_order);
93 return Some(*nth);
94 }
95
96 let (lower, upper_mid, _) = values.select_nth_unstable_by(n / 2, total_order);
100 let upper = *upper_mid;
101 let lower_mid = lower.iter().copied().reduce(|a, b| {
102 if total_order(&b, &a) == Ordering::Greater {
103 b
104 } else {
105 a
106 }
107 })?;
108 let two = A::from(2.0)?;
109 Some((lower_mid + upper) / two)
110}
111
112pub fn median<A: Float>(values: &[A]) -> Option<A> {
115 let mut scratch: Vec<A> = values.to_vec();
116 median_in_place(&mut scratch)
117}
118
119pub fn standard_normal_sf(z: f64) -> Result<f64, String> {
124 let dist = scirs2_stats::distributions::normal::Normal::new(0.0_f64, 1.0_f64)
125 .map_err(|e| format!("standard normal construction failed: {e}"))?;
126 Ok((1.0 - dist.cdf(z)).clamp(0.0, 1.0))
127}
128
129pub fn normal_two_sided_p(z: f64) -> Result<f64, String> {
131 let upper = standard_normal_sf(z.abs())?;
132 Ok((2.0 * upper).clamp(0.0, 1.0))
133}
134
135fn ln_gamma(x: f64) -> f64 {
139 const COEFFICIENTS: [f64; 9] = [
140 0.999_999_999_999_810,
141 676.520_368_121_885,
142 -1_259.139_216_722_403,
143 771.323_428_777_653,
144 -176.615_029_162_141,
145 12.507_343_278_687,
146 -0.138_571_095_265_720,
147 0.000_009_984_369_578,
148 0.000_000_150_563_274,
149 ];
150
151 if x < 0.5 {
152 return (std::f64::consts::PI / (std::f64::consts::PI * x).sin()).ln() - ln_gamma(1.0 - x);
154 }
155
156 let x = x - 1.0;
157 let mut series = COEFFICIENTS[0];
158 for (index, coefficient) in COEFFICIENTS.iter().enumerate().skip(1) {
159 series += coefficient / (x + index as f64);
160 }
161 let t = x + 7.5;
162 0.5 * (2.0 * std::f64::consts::PI).ln() + (x + 0.5) * t.ln() - t + series.ln()
163}
164
165fn regularized_upper_gamma(a: f64, x: f64) -> Result<f64, String> {
171 if !a.is_finite() || a <= 0.0 || !x.is_finite() || x < 0.0 {
172 return Err(format!(
173 "regularized_upper_gamma requires a > 0 and x >= 0, got a={a}, x={x}"
174 ));
175 }
176 if x == 0.0 {
177 return Ok(1.0);
178 }
179
180 let log_prefactor = -x + a * x.ln() - ln_gamma(a);
181
182 if x < a + 1.0 {
183 let mut term = 1.0 / a;
185 let mut sum = term;
186 let mut n = a;
187 for _ in 0..1000 {
188 n += 1.0;
189 term *= x / n;
190 sum += term;
191 if term.abs() < sum.abs() * 1e-16 {
192 break;
193 }
194 }
195 let lower = sum * log_prefactor.exp();
196 return Ok((1.0 - lower).clamp(0.0, 1.0));
197 }
198
199 let tiny = 1e-300_f64;
202 let mut b = x + 1.0 - a;
203 let mut c = 1.0 / tiny;
204 let mut d = 1.0 / b;
205 let mut h = d;
206 for i in 1..1000 {
207 let an = -(i as f64) * (i as f64 - a);
208 b += 2.0;
209 d = an * d + b;
210 if d.abs() < tiny {
211 d = tiny;
212 }
213 c = b + an / c;
214 if c.abs() < tiny {
215 c = tiny;
216 }
217 d = 1.0 / d;
218 let delta = d * c;
219 h *= delta;
220 if (delta - 1.0).abs() < 1e-16 {
221 break;
222 }
223 }
224
225 Ok((log_prefactor.exp() * h).clamp(0.0, 1.0))
226}
227
228pub fn chi_square_sf(x: f64, df: f64) -> Result<f64, String> {
238 if !(df.is_finite() && df > 0.0) {
239 return Err(format!("chi-square requires df > 0, got {df}"));
240 }
241 if x <= 0.0 {
242 return Ok(1.0);
243 }
244 regularized_upper_gamma(df / 2.0, x / 2.0)
245}
246
247pub fn kolmogorov_sf(z: f64) -> f64 {
255 if !z.is_finite() || z <= 0.0 {
256 return 1.0;
257 }
258
259 let a2 = -2.0 * z * z;
260 let mut sign = 2.0_f64;
261 let mut sum = 0.0_f64;
262 let mut previous_magnitude = 0.0_f64;
263
264 for k in 1..=200_u32 {
265 let term = sign * (a2 * f64::from(k * k)).exp();
266 sum += term;
267 let magnitude = term.abs();
268 if magnitude <= 1e-8 * previous_magnitude || magnitude <= 1e-16 * sum.abs() {
269 break;
270 }
271 previous_magnitude = magnitude;
272 sign = -sign;
273 }
274
275 sum.clamp(0.0, 1.0)
276}
277
278pub fn ks_two_sample_p(d: f64, n1: usize, n2: usize) -> f64 {
281 if n1 == 0 || n2 == 0 {
282 return 1.0;
283 }
284 let n1 = n1 as f64;
285 let n2 = n2 as f64;
286 let effective_n = (n1 * n2) / (n1 + n2);
287 kolmogorov_sf(effective_n.sqrt() * d)
288}
289
290pub fn ks_statistic<A: Float>(sample_a: &[A], sample_b: &[A]) -> Option<f64> {
296 if sample_a.is_empty() || sample_b.is_empty() {
297 return None;
298 }
299
300 let mut a: Vec<A> = sample_a.to_vec();
301 let mut b: Vec<A> = sample_b.to_vec();
302 sort_ascending(&mut a);
303 sort_ascending(&mut b);
304
305 let na = a.len();
306 let nb = b.len();
307 let mut i = 0usize;
308 let mut j = 0usize;
309 let mut max_diff = 0.0_f64;
310
311 while i < na || j < nb {
312 let next = match (a.get(i), b.get(j)) {
319 (Some(x), Some(y)) => {
320 if total_order(x, y) == Ordering::Greater {
321 *y
322 } else {
323 *x
324 }
325 }
326 (Some(x), None) => *x,
327 (None, Some(y)) => *y,
328 (None, None) => break,
329 };
330
331 while i < na && total_order(&a[i], &next) != Ordering::Greater {
332 i += 1;
333 }
334 while j < nb && total_order(&b[j], &next) != Ordering::Greater {
335 j += 1;
336 }
337
338 let ecdf_a = i as f64 / na as f64;
339 let ecdf_b = j as f64 / nb as f64;
340 let diff = (ecdf_a - ecdf_b).abs();
341 if diff > max_diff {
342 max_diff = diff;
343 }
344 }
345
346 Some(max_diff)
347}
348
349pub fn finite_range<A: Float>(values: &[A]) -> Option<(f64, f64)> {
352 let mut min = f64::INFINITY;
353 let mut max = f64::NEG_INFINITY;
354 for value in values {
355 let Some(v) = value.to_f64() else { continue };
356 if !v.is_finite() {
357 continue;
358 }
359 if v < min {
360 min = v;
361 }
362 if v > max {
363 max = v;
364 }
365 }
366 if min.is_finite() && max.is_finite() {
367 Some((min, max))
368 } else {
369 None
370 }
371}
372
373pub fn histogram_counts<A: Float>(values: &[A], min: f64, max: f64, bins: usize) -> Vec<f64> {
379 let bins = bins.max(1);
380 let mut counts = vec![0.0_f64; bins];
381 let width = if max > min {
382 (max - min) / bins as f64
383 } else {
384 0.0
387 };
388
389 for value in values {
390 let Some(v) = value.to_f64() else { continue };
391 if !v.is_finite() {
392 continue;
393 }
394 let index = if width > 0.0 {
395 (((v - min) / width).floor().max(0.0) as usize).min(bins - 1)
396 } else {
397 0
398 };
399 counts[index] += 1.0;
400 }
401
402 counts
403}
404
405pub fn smoothed_pmf(counts: &[f64], smoothing: f64) -> Vec<f64> {
411 let smoothing = smoothing.max(f64::MIN_POSITIVE);
412 let total: f64 = counts.iter().sum::<f64>() + smoothing * counts.len() as f64;
413 if total <= 0.0 {
414 let uniform = 1.0 / counts.len().max(1) as f64;
415 return vec![uniform; counts.len()];
416 }
417 counts.iter().map(|&c| (c + smoothing) / total).collect()
418}
419
420pub fn kl_divergence(p: &[f64], q: &[f64]) -> Result<f64, String> {
423 if p.len() != q.len() {
424 return Err("KL divergence requires equal-length distributions".to_string());
425 }
426 let mut sum = 0.0_f64;
427 for (&pi, &qi) in p.iter().zip(q.iter()) {
428 if pi <= 0.0 {
429 continue;
430 }
431 if qi <= 0.0 {
432 return Err("KL divergence is undefined for a zero reference bin".to_string());
433 }
434 sum += pi * (pi / qi).ln();
435 }
436 Ok(sum.max(0.0))
437}
438
439pub fn js_divergence(p: &[f64], q: &[f64]) -> Result<f64, String> {
441 if p.len() != q.len() {
442 return Err("JS divergence requires equal-length distributions".to_string());
443 }
444 let mixture: Vec<f64> = p
445 .iter()
446 .zip(q.iter())
447 .map(|(&pi, &qi)| 0.5 * (pi + qi))
448 .collect();
449 let left = kl_divergence(p, &mixture)?;
450 let right = kl_divergence(q, &mixture)?;
451 Ok((0.5 * left + 0.5 * right).clamp(0.0, std::f64::consts::LN_2))
452}
453
454pub fn hellinger_distance(p: &[f64], q: &[f64]) -> Result<f64, String> {
456 if p.len() != q.len() {
457 return Err("Hellinger distance requires equal-length distributions".to_string());
458 }
459 let bhattacharyya: f64 = p
460 .iter()
461 .zip(q.iter())
462 .map(|(&pi, &qi)| (pi.max(0.0) * qi.max(0.0)).sqrt())
463 .sum();
464 Ok((1.0 - bhattacharyya).max(0.0).sqrt())
465}
466
467pub fn wasserstein_1d<A: Float>(sample_a: &[A], sample_b: &[A]) -> Option<f64> {
475 if sample_a.is_empty() || sample_b.is_empty() {
476 return None;
477 }
478
479 let mut a: Vec<f64> = sample_a.iter().filter_map(|v| v.to_f64()).collect();
480 let mut b: Vec<f64> = sample_b.iter().filter_map(|v| v.to_f64()).collect();
481 if a.is_empty() || b.is_empty() {
482 return None;
483 }
484 a.sort_by(|x, y| x.partial_cmp(y).unwrap_or(Ordering::Equal));
485 b.sort_by(|x, y| x.partial_cmp(y).unwrap_or(Ordering::Equal));
486
487 let na = a.len();
488 let nb = b.len();
489 let mut i = 0usize;
490 let mut j = 0usize;
491 let mut previous = a[0].min(b[0]);
492 let mut total = 0.0_f64;
493
494 while i < na || j < nb {
495 let next = match (a.get(i), b.get(j)) {
496 (Some(&x), Some(&y)) => x.min(y),
497 (Some(&x), None) => x,
498 (None, Some(&y)) => y,
499 (None, None) => break,
500 };
501
502 let ecdf_a = i as f64 / na as f64;
505 let ecdf_b = j as f64 / nb as f64;
506 total += (ecdf_a - ecdf_b).abs() * (next - previous);
507 previous = next;
508
509 while i < na && a[i] <= next {
510 i += 1;
511 }
512 while j < nb && b[j] <= next {
513 j += 1;
514 }
515 }
516
517 Some(total)
518}
519
520pub fn g_test_statistic(counts_a: &[f64], counts_b: &[f64]) -> Result<(f64, f64), String> {
528 if counts_a.len() != counts_b.len() {
529 return Err("G-test requires equal-length histograms".to_string());
530 }
531 let total_a: f64 = counts_a.iter().sum();
532 let total_b: f64 = counts_b.iter().sum();
533 let grand_total = total_a + total_b;
534 if total_a <= 0.0 || total_b <= 0.0 {
535 return Err("G-test requires both samples to be non-empty".to_string());
536 }
537
538 let mut g = 0.0_f64;
539 let mut non_empty_columns = 0usize;
540 for (&observed_a, &observed_b) in counts_a.iter().zip(counts_b.iter()) {
541 let column_total = observed_a + observed_b;
542 if column_total <= 0.0 {
543 continue;
544 }
545 non_empty_columns += 1;
546
547 let expected_a = column_total * total_a / grand_total;
548 let expected_b = column_total * total_b / grand_total;
549 if observed_a > 0.0 && expected_a > 0.0 {
550 g += observed_a * (observed_a / expected_a).ln();
551 }
552 if observed_b > 0.0 && expected_b > 0.0 {
553 g += observed_b * (observed_b / expected_b).ln();
554 }
555 }
556
557 let degrees_of_freedom = (non_empty_columns.saturating_sub(1)) as f64;
558 Ok((2.0 * g, degrees_of_freedom))
559}
560
561#[cfg(test)]
562mod tests {
563 use super::*;
564
565 #[test]
566 fn total_order_is_nan_safe_and_total() {
567 let mut values = vec![3.0_f64, f64::NAN, 1.0, 2.0];
568 sort_ascending(&mut values);
570 assert_eq!(values[0], 1.0);
571 assert_eq!(values[1], 2.0);
572 assert_eq!(values[2], 3.0);
573 assert!(values[3].is_nan(), "NaN must sort last");
574 }
575
576 #[test]
577 fn median_matches_definition_for_odd_and_even_lengths() {
578 assert_eq!(median(&[5.0_f64, 1.0, 3.0]), Some(3.0));
579 assert_eq!(median(&[4.0_f64, 1.0, 3.0, 2.0]), Some(2.5));
580 assert_eq!(median::<f64>(&[]), None);
581 }
582
583 #[test]
584 fn quantile_uses_nearest_rank() {
585 let mut values = vec![1.0_f64, 2.0, 3.0, 4.0];
586 assert_eq!(quantile_in_place(&mut values, 0.0), Some(1.0));
587 let mut values = vec![1.0_f64, 2.0, 3.0, 4.0];
588 assert_eq!(quantile_in_place(&mut values, 1.0), Some(4.0));
589 let mut values = vec![1.0_f64, 2.0, 3.0, 4.0];
590 assert_eq!(quantile_in_place(&mut values, 0.5), Some(2.0));
591 }
592
593 #[test]
594 fn normal_survival_function_matches_known_quantiles() {
595 let p = standard_normal_sf(1.959_963_985).expect("normal sf");
596 assert!(
597 (p - 0.025).abs() < 1e-4,
598 "P(Z > 1.96) should be ~0.025, got {p}"
599 );
600 let p0 = standard_normal_sf(0.0).expect("normal sf");
601 assert!((p0 - 0.5).abs() < 1e-9);
602 }
603
604 #[test]
606 fn chi_square_survival_function_matches_known_quantiles() {
607 let p = chi_square_sf(3.841_458_8, 1.0).expect("chi2 sf");
609 assert!((p - 0.05).abs() < 1e-6, "expected 0.05, got {p}");
610
611 let p = chi_square_sf(5.991_464_5, 2.0).expect("chi2 sf");
613 assert!((p - 0.05).abs() < 1e-6, "expected 0.05, got {p}");
614
615 let p = chi_square_sf(24.995_79, 15.0).expect("chi2 sf");
617 assert!((p - 0.05).abs() < 1e-5, "expected 0.05, got {p}");
618 }
619
620 #[test]
624 fn chi_square_survival_function_is_correct_in_the_series_branch() {
625 let p = chi_square_sf(1.0, 4.0).expect("chi2 sf");
627 assert!((p - 0.909_796).abs() < 1e-5, "expected 0.909796, got {p}");
628
629 let p = chi_square_sf(2.0, 10.0).expect("chi2 sf");
631 assert!((p - 0.996_340).abs() < 1e-5, "expected 0.996340, got {p}");
632
633 assert!(chi_square_sf(1e-12, 3.0).expect("chi2 sf") > 0.999_999);
635 assert_eq!(chi_square_sf(0.0, 3.0).expect("chi2 sf"), 1.0);
636 }
637
638 #[test]
641 fn chi_square_survival_function_is_continuous_across_the_branch_switch() {
642 let df = 6.0_f64;
643 let a = df / 2.0; let switch = 2.0 * (a + 1.0);
645 let below = chi_square_sf(switch - 1e-7, df).expect("chi2 sf");
646 let above = chi_square_sf(switch + 1e-7, df).expect("chi2 sf");
647 assert!(
653 (below - above).abs() < 1e-7,
654 "the two branches disagree at the switch point ({below} vs {above})"
655 );
656
657 let mut previous = 1.0_f64;
658 for step in 1..=200 {
659 let x = step as f64 * 0.15;
660 let p = chi_square_sf(x, df).expect("chi2 sf");
661 assert!(
662 p <= previous + 1e-12,
663 "survival function increased at x = {x} ({previous} -> {p})"
664 );
665 previous = p;
666 }
667 }
668
669 #[test]
670 fn ln_gamma_matches_known_values() {
671 assert!(ln_gamma(1.0).abs() < 1e-12);
673 assert!(ln_gamma(2.0).abs() < 1e-12);
674 assert!((ln_gamma(5.0) - 24.0_f64.ln()).abs() < 1e-11);
675 assert!(
676 (ln_gamma(0.5) - std::f64::consts::PI.sqrt().ln()).abs() < 1e-11,
677 "ln G(0.5) = {}",
678 ln_gamma(0.5)
679 );
680 }
681
682 #[test]
683 fn ks_statistic_is_one_for_disjoint_samples() {
684 let d = ks_statistic(&[0.0_f64, 1.0, 2.0], &[10.0, 11.0, 12.0]).expect("ks");
685 assert!(
686 (d - 1.0).abs() < 1e-12,
687 "disjoint samples must give D = 1, got {d}"
688 );
689 let p = ks_two_sample_p(d, 3, 3);
690 assert!(
691 p < 0.5,
692 "D = 1 on n = 3 should be at least mildly significant, got {p}"
693 );
694 }
695
696 #[test]
697 fn ks_statistic_is_zero_for_identical_samples() {
698 let d = ks_statistic(&[1.0_f64, 2.0, 3.0], &[1.0, 2.0, 3.0]).expect("ks");
699 assert!(
700 d.abs() < 1e-12,
701 "identical samples must give D = 0, got {d}"
702 );
703 assert!((ks_two_sample_p(d, 3, 3) - 1.0).abs() < 1e-12);
704 }
705
706 #[test]
707 fn wasserstein_matches_closed_form_for_a_pure_shift() {
708 let a = [0.0_f64, 1.0, 2.0, 3.0];
710 let b = [5.0_f64, 6.0, 7.0, 8.0];
711 let w = wasserstein_1d(&a, &b).expect("w1");
712 assert!((w - 5.0).abs() < 1e-9, "expected W1 = 5, got {w}");
713 }
714
715 #[test]
716 fn divergences_are_zero_for_identical_distributions_and_positive_otherwise() {
717 let p = smoothed_pmf(&[10.0, 10.0, 10.0], 0.5);
718 let q = smoothed_pmf(&[10.0, 10.0, 10.0], 0.5);
719 assert!(kl_divergence(&p, &q).expect("kl") < 1e-12);
720 assert!(js_divergence(&p, &q).expect("js") < 1e-12);
721 assert!(hellinger_distance(&p, &q).expect("hellinger") < 1e-6);
722
723 let r = smoothed_pmf(&[30.0, 0.0, 0.0], 0.5);
724 assert!(kl_divergence(&p, &r).expect("kl") > 0.1);
725 assert!(js_divergence(&p, &r).expect("js") > 0.1);
726 assert!(hellinger_distance(&p, &r).expect("hellinger") > 0.1);
727 }
728
729 #[test]
730 fn g_test_is_insignificant_for_matching_histograms() {
731 let (g, df) = g_test_statistic(&[20.0, 20.0, 20.0], &[20.0, 20.0, 20.0]).expect("g-test");
732 assert!(g.abs() < 1e-9);
733 assert_eq!(df, 2.0);
734 let p = chi_square_sf(g, df.max(1.0)).expect("chi2");
735 assert!(
736 p > 0.9,
737 "identical histograms must not be significant, got {p}"
738 );
739 }
740
741 #[test]
742 fn g_test_is_significant_for_disjoint_histograms() {
743 let (g, df) = g_test_statistic(&[60.0, 0.0], &[0.0, 60.0]).expect("g-test");
744 assert!(
745 g > 100.0,
746 "disjoint histograms should give a large G, got {g}"
747 );
748 let p = chi_square_sf(g, df.max(1.0)).expect("chi2");
749 assert!(p < 1e-6, "expected an extremely small p-value, got {p}");
750 }
751}