1use crate::aberth::{aberth, aberth_atomic, aberth_autocorr, aberth_mt};
9use crate::autocorr::{pbairstow_autocorr, pbairstow_autocorr_atomic, pbairstow_autocorr_mt};
10use crate::execution_policy::should_parallelize;
11use crate::rootfinding::{pbairstow_even, pbairstow_even_atomic, pbairstow_even_mt, Options};
12use num_complex::Complex;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum SolveMode {
17 Automatic,
19 Sequential,
20 MultiThreaded,
21 Atomic,
22}
23
24pub fn solve_aberth(
29 coeffs: &[f64],
30 zs: &mut [Complex<f64>],
31 options: &Options,
32 mode: SolveMode,
33) -> (usize, bool) {
34 match mode {
35 SolveMode::Sequential => aberth(coeffs, zs, options),
36 SolveMode::MultiThreaded => aberth_mt(coeffs, zs, options),
37 SolveMode::Atomic => aberth_atomic(coeffs, zs, options),
38 SolveMode::Automatic => {
39 if should_parallelize(zs.len()) {
40 aberth_mt(coeffs, zs, options)
41 } else {
42 aberth(coeffs, zs, options)
43 }
44 }
45 }
46}
47
48pub fn solve_pbairstow_even(
50 coeffs: &[f64],
51 vrs: &mut [crate::rootfinding::Vec2],
52 options: &Options,
53 mode: SolveMode,
54) -> (usize, bool) {
55 match mode {
56 SolveMode::Sequential => pbairstow_even(coeffs, vrs, options),
57 SolveMode::MultiThreaded => pbairstow_even_mt(coeffs, vrs, options),
58 SolveMode::Atomic => pbairstow_even_atomic(coeffs, vrs, options),
59 SolveMode::Automatic => {
60 if should_parallelize(vrs.len()) {
61 pbairstow_even_mt(coeffs, vrs, options)
62 } else {
63 pbairstow_even(coeffs, vrs, options)
64 }
65 }
66 }
67}
68
69pub fn solve_pbairstow_autocorr(
71 coeffs: &[f64],
72 vrs: &mut [crate::rootfinding::Vec2],
73 options: &Options,
74 mode: SolveMode,
75) -> (usize, bool) {
76 match mode {
77 SolveMode::Sequential => pbairstow_autocorr(coeffs, vrs, options),
78 SolveMode::MultiThreaded => pbairstow_autocorr_mt(coeffs, vrs, options),
79 SolveMode::Atomic => pbairstow_autocorr_atomic(coeffs, vrs, options),
80 SolveMode::Automatic => {
81 if should_parallelize(vrs.len()) {
82 pbairstow_autocorr_mt(coeffs, vrs, options)
83 } else {
84 pbairstow_autocorr(coeffs, vrs, options)
85 }
86 }
87 }
88}
89
90#[inline]
95pub fn solve_aberth_autocorr(
96 coeffs: &[f64],
97 zs: &mut [Complex<f64>],
98 options: &Options,
99 _mode: SolveMode,
100) -> (usize, bool) {
101 aberth_autocorr(coeffs, zs, options)
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107
108 #[test]
109 fn test_should_parallelize_threshold() {
110 assert!(!should_parallelize(0));
111 assert!(!should_parallelize(4));
112 assert!(should_parallelize(5));
113 }
114
115 #[test]
116 fn test_solve_aberth_automatic_small() {
117 let coeffs = vec![10.0, 34.0, 75.0, 94.0, 150.0, 94.0, 75.0, 34.0, 10.0];
119 let mut zrs = crate::aberth::initial_aberth(&coeffs);
120 let (_, found) = solve_aberth(&coeffs, &mut zrs, &Options::default(), SolveMode::Automatic);
121 assert!(found);
122 }
123
124 #[test]
125 fn test_solve_aberth_explicit_modes() {
126 let coeffs = vec![
127 1.0, 0.0, 2.0, 0.0, 3.0, 0.0, 4.0, 0.0, 3.0, 0.0, 2.0, 0.0, 1.0,
128 ];
129 let mut zrs = crate::aberth::initial_aberth(&coeffs);
130 let (_, found) = solve_aberth(
131 &coeffs,
132 &mut zrs,
133 &Options::default(),
134 SolveMode::MultiThreaded,
135 );
136 assert!(found);
137 let mut zrs = crate::aberth::initial_aberth(&coeffs);
138 let (_, found) = solve_aberth(&coeffs, &mut zrs, &Options::default(), SolveMode::Atomic);
139 assert!(found);
140 }
141
142 #[test]
143 fn test_solve_pbairstow_even() {
144 let coeffs = vec![10.0, 34.0, 75.0, 94.0, 150.0, 94.0, 75.0, 34.0, 10.0];
145 let mut vrs = crate::rootfinding::initial_guess(&coeffs);
146 let (_, found) =
147 solve_pbairstow_even(&coeffs, &mut vrs, &Options::default(), SolveMode::Automatic);
148 assert!(found);
149 }
150}